diff --git a/aixplain/modules/agent/__init__.py b/aixplain/modules/agent/__init__.py index 69c22f03..33d0a7fb 100644 --- a/aixplain/modules/agent/__init__.py +++ b/aixplain/modules/agent/__init__.py @@ -306,6 +306,7 @@ def run( max_iterations: int = 5, output_format: Optional[OutputFormat] = None, expected_output: Optional[Union[BaseModel, Text, dict]] = None, + trace_request: bool = False, ) -> AgentResponse: """Runs an agent call. @@ -323,7 +324,7 @@ def run( max_iterations (int, optional): maximum number of iterations between the agent and the tools. Defaults to 10. output_format (OutputFormat, optional): response format. If not provided, uses the format set during initialization. expected_output (Union[BaseModel, Text, dict], optional): expected output. Defaults to None. - + trace_request (bool, optional): return the request id for tracing the request. Defaults to False. Returns: Dict: parsed output from model """ @@ -352,6 +353,7 @@ def run( max_iterations=max_iterations, output_format=output_format, expected_output=expected_output, + trace_request=trace_request, ) if response["status"] == ResponseStatus.FAILED: end = time.time() @@ -406,6 +408,7 @@ def run_async( output_format: Optional[OutputFormat] = None, expected_output: Optional[Union[BaseModel, Text, dict]] = None, evolve: Union[Dict[str, Any], EvolveParam, None] = None, + trace_request: bool = False, ) -> AgentResponse: """Runs asynchronously an agent call. @@ -423,7 +426,7 @@ def run_async( expected_output (Union[BaseModel, Text, dict], optional): expected output. Defaults to None. output_format (ResponseFormat, optional): response format. Defaults to TEXT. evolve (Union[Dict[str, Any], EvolveParam, None], optional): evolve the agent configuration. Can be a dictionary, EvolveParam instance, or None. - + trace_request (bool, optional): return the request id for tracing the request. Defaults to False. Returns: dict: polling URL in response """ @@ -516,6 +519,8 @@ def run_async( try: r = _request_with_retry("post", self.url, headers=headers, data=payload) resp = r.json() + if trace_request: + logging.info(f"Agent Run Async: Trace request id: {resp.get('requestId')}") poll_url = resp.get("data") return AgentResponse( status=ResponseStatus.IN_PROGRESS, diff --git a/aixplain/modules/team_agent/__init__.py b/aixplain/modules/team_agent/__init__.py index 74b9a14c..58c19cd4 100644 --- a/aixplain/modules/team_agent/__init__.py +++ b/aixplain/modules/team_agent/__init__.py @@ -200,6 +200,7 @@ def run( max_iterations: int = 30, output_format: Optional[OutputFormat] = None, expected_output: Optional[Union[BaseModel, Text, dict]] = None, + trace_request: bool = False, ) -> AgentResponse: """Runs a team agent call. @@ -217,6 +218,7 @@ def run( max_iterations (int, optional): maximum number of iterations between the agents. Defaults to 30. output_format (OutputFormat, optional): response format. If not provided, uses the format set during initialization. expected_output (Union[BaseModel, Text, dict], optional): expected output. Defaults to None. + trace_request (bool, optional): return the request id for tracing the request. Defaults to False. Returns: AgentResponse: parsed output from model """ @@ -243,6 +245,7 @@ def run( max_iterations=max_iterations, output_format=output_format, expected_output=expected_output, + trace_request=trace_request, ) if response["status"] == ResponseStatus.FAILED: end = time.time() @@ -289,6 +292,7 @@ def run_async( output_format: Optional[OutputFormat] = None, expected_output: Optional[Union[BaseModel, Text, dict]] = None, evolve: Union[Dict[str, Any], EvolveParam, None] = None, + trace_request: bool = False, ) -> AgentResponse: """Runs asynchronously a Team Agent call. @@ -305,6 +309,7 @@ def run_async( output_format (OutputFormat, optional): response format. If not provided, uses the format set during initialization. expected_output (Union[BaseModel, Text, dict], optional): expected output. Defaults to None. evolve (Union[Dict[str, Any], EvolveParam, None], optional): evolve the team agent configuration. Can be a dictionary, EvolveParam instance, or None. + trace_request (bool, optional): return the request id for tracing the request. Defaults to False. Returns: AgentResponse: polling URL in response """ @@ -393,7 +398,8 @@ def run_async( try: resp = r.json() logging.info(f"Result of request for {name} - {r.status_code} - {resp}") - + if trace_request: + logging.info(f"Team Agent Run Async: Trace request id: {resp.get('requestId')}") poll_url = resp["data"] response = AgentResponse( status=ResponseStatus.IN_PROGRESS, diff --git a/pyproject.toml b/pyproject.toml index bfb57959..c9022545 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ namespaces = true [project] name = "aiXplain" -version = "0.2.35" +version = "0.2.36rc" description = "aiXplain SDK adds AI functions to software." readme = "README.md" requires-python = ">=3.9, <4"