agno 1.7.5__py3-none-any.whl → 1.7.7__py3-none-any.whl

agno/models/portkey/portkey.py

@@ -0,0 +1,88 @@
+from dataclasses import dataclass
+from os import getenv
+from typing import Any, Dict, Optional, cast
+
+from agno.exceptions import ModelProviderError
+from agno.models.openai.like import OpenAILike
+
+try:
+    from portkey_ai import PORTKEY_GATEWAY_URL, createHeaders
+except ImportError:
+    raise ImportError("`portkey-ai` not installed. Please install using `pip install portkey-ai`")
+
+
+@dataclass
+class Portkey(OpenAILike):
+    """
+    A class for using models through the Portkey AI Gateway.
+
+    Attributes:
+        id (str): The model id. Defaults to "gpt-4o-mini".
+        name (str): The model name. Defaults to "Portkey".
+        provider (str): The provider name. Defaults to "Portkey".
+        portkey_api_key (Optional[str]): The Portkey API key.
+        virtual_key (Optional[str]): The virtual key for model routing.
+        config (Optional[Dict[str, Any]]): Portkey configuration for routing, retries, etc.
+        base_url (str): The Portkey gateway URL.
+    """
+
+    id: str = "gpt-4o-mini"
+    name: str = "Portkey"
+    provider: str = "Portkey"
+
+    portkey_api_key: Optional[str] = getenv("PORTKEY_API_KEY")
+    virtual_key: Optional[str] = getenv("PORTKEY_VIRTUAL_KEY")
+    config: Optional[Dict[str, Any]] = None
+    base_url: str = PORTKEY_GATEWAY_URL
+
+    def _get_client_params(self) -> Dict[str, Any]:
+        # Check for required keys
+        if not self.portkey_api_key:
+            raise ModelProviderError(
+                message="PORTKEY_API_KEY not set. Please set the PORTKEY_API_KEY environment variable.",
+                model_name=self.name,
+                model_id=self.id,
+            )
+
+        if not self.virtual_key:
+            raise ModelProviderError(
+                message="PORTKEY_VIRTUAL_KEY not set. Please set the PORTKEY_VIRTUAL_KEY environment variable.",
+                model_name=self.name,
+                model_id=self.id,
+            )
+
+        # Create headers using Portkey's createHeaders function
+        header_params: Dict[str, Any] = {
+            "api_key": self.portkey_api_key,
+            "virtual_key": self.virtual_key,
+        }
+
+        if self.config is not None:
+            header_params["config"] = self.config
+
+        portkey_headers = cast(Dict[str, Any], createHeaders(**header_params))
+
+        # Merge with any existing default headers
+        default_headers: Dict[str, Any] = {}
+        if self.default_headers and isinstance(self.default_headers, dict):
+            default_headers.update(self.default_headers)
+        default_headers.update(portkey_headers)
+
+        # Define base client params
+        base_params = {
+            "api_key": "not-needed",  # We use virtual keys instead
+            "organization": self.organization,
+            "base_url": self.base_url,
+            "timeout": self.timeout,
+            "max_retries": self.max_retries,
+            "default_headers": default_headers,
+            "default_query": self.default_query,
+        }
+
+        # Create client_params dict with non-None values
+        client_params = {k: v for k, v in base_params.items() if v is not None}
+
+        # Add additional client params if provided
+        if self.client_params:
+            client_params.update(self.client_params)
+        return client_params

agno/models/xai/xai.py

@@ -4,9 +4,17 @@ from typing import Any, Dict, List, Optional, Type, Union
 
 from pydantic import BaseModel
 
+from agno.models.message import Citations, UrlCitation
 from agno.models.openai.like import OpenAILike
+from agno.models.response import ModelResponse
 from agno.utils.log import log_debug
 
+try:
+    from openai.types.chat.chat_completion import ChatCompletion
+    from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+except (ImportError, ModuleNotFoundError):
+    raise ImportError("`openai` not installed. Please install using `pip install openai`")
+
 
 @dataclass
 class xAI(OpenAILike):
@@ -56,3 +64,49 @@ class xAI(OpenAILike):
             log_debug(f"Calling {self.provider} with request parameters: {request_params}", log_level=2)
 
         return request_params
+
+    def parse_provider_response(
+        self,
+        response: ChatCompletion,
+        response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
+    ) -> ModelResponse:
+        """
+        Parse the xAI response into a ModelResponse.
+        """
+        model_response = super().parse_provider_response(response, response_format)
+
+        if hasattr(response, "citations") and response.citations:
+            citations = Citations()
+            url_citations = []
+            for citation_url in response.citations:
+                url_citations.append(UrlCitation(url=str(citation_url)))
+
+            citations.urls = url_citations
+            citations.raw = response.citations
+            model_response.citations = citations
+
+        return model_response
+
+    def parse_provider_response_delta(self, response_delta: ChatCompletionChunk) -> ModelResponse:
+        """
+        Parse the xAI streaming response.
+
+        Args:
+            response_delta: Raw response chunk
+
+        Returns:
+            ModelResponse: Parsed response data
+        """
+        model_response = super().parse_provider_response_delta(response_delta)
+
+        if hasattr(response_delta, "citations") and response_delta.citations:
+            citations = Citations()
+            url_citations = []
+            for citation_url in response_delta.citations:
+                url_citations.append(UrlCitation(url=str(citation_url)))
+
+            citations.urls = url_citations
+            citations.raw = response_delta.citations
+            model_response.citations = citations
+
+        return model_response

agno/run/v2/workflow.py

@@ -561,3 +561,7 @@ class WorkflowRunResponse:
             return self.content.model_dump_json(exclude_none=True, **kwargs)
         else:
             return json.dumps(self.content, **kwargs)
+
+    def has_completed(self) -> bool:
+        """Check if the workflow run is completed (either successfully or with error)"""
+        return self.status in [RunStatus.completed, RunStatus.error]

agno/storage/mysql.py

@@ -131,7 +131,9 @@ class MySQLStorage(Storage):
         elif self.mode == "workflow_v2":
             specific_columns = [
                 Column("workflow_id", String(255), index=True),
+                Column("workflow_name", String(255), index=True),
                 Column("workflow_data", JSON),
+                Column("runs", JSON),
             ]
         # Create table with all columns
         table = Table(

agno/storage/postgres.py

@@ -30,7 +30,7 @@ class PostgresStorage(Storage):
         db_engine: Optional[Engine] = None,
         schema_version: int = 1,
         auto_upgrade_schema: bool = False,
-        mode: Optional[Literal["agent", "team", "workflow"]] = "agent",
+        mode: Optional[Literal["agent", "team", "workflow", "workflow_v2"]] = "agent",
     ):
         """
         This class provides agent storage using a PostgreSQL table.
@@ -131,7 +131,9 @@ class PostgresStorage(Storage):
         elif self.mode == "workflow_v2":
             specific_columns = [
                 Column("workflow_id", String, index=True),
+                Column("workflow_name", String, index=True),
                 Column("workflow_data", postgresql.JSONB),
+                Column("runs", postgresql.JSONB),
             ]
 
         # Create table with all columns
@@ -569,9 +571,9 @@ class PostgresStorage(Storage):
                     stmt = postgresql.insert(self.table).values(
                         session_id=session.session_id,
                         workflow_id=session.workflow_id,  # type: ignore
-                        workflow_name=session.workflow_name,  # type: ignore
                         user_id=session.user_id,
                         runs=session_dict.get("runs"),
+                        workflow_name=session.workflow_name,  # type: ignore
                         workflow_data=session.workflow_data,  # type: ignore
                         session_data=session.session_data,
                         extra_data=session.extra_data,
@@ -582,9 +584,9 @@ class PostgresStorage(Storage):
                         index_elements=["session_id"],
                         set_=dict(
                             workflow_id=session.workflow_id,  # type: ignore
-                            workflow_name=session.workflow_name,  # type: ignore
                             user_id=session.user_id,
                             runs=session_dict.get("runs"),
+                            workflow_name=session.workflow_name,  # type: ignore
                             workflow_data=session.workflow_data,  # type: ignore
                             session_data=session.session_data,
                             extra_data=session.extra_data,

agno/storage/session/v2/workflow.py

@@ -40,21 +40,45 @@ class WorkflowSession:
         if self.runs is None:
             self.runs = []
 
-    def add_run(self, run_response: WorkflowRunResponse) -> None:
-        """Add a workflow run response to this session"""
+    def upsert_run(self, run: WorkflowRunResponse) -> None:
+        """Add or update a workflow run (upsert behavior)"""
         if self.runs is None:
             self.runs = []
-        # Store the actual WorkflowRunResponse object
-        self.runs.append(run_response)
+
+        # Find existing run and update it, or append new one
+        for i, existing_run in enumerate(self.runs):
+            if existing_run.run_id == run.run_id:
+                # Update existing run
+                self.runs[i] = run
+                return
+
+        # Run not found, append new one
+        self.runs.append(run)
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary for storage, serializing runs to dicts"""
+
+        runs_data = None
+        if self.runs:
+            runs_data = []
+            for run in self.runs:
+                try:
+                    runs_data.append(run.to_dict())
+                except Exception as e:
+                    # If run serialization fails, create a minimal representation
+                    runs_data.append(
+                        {
+                            "run_id": getattr(run, "run_id", "unknown"),
+                            "status": str(getattr(run, "status", "unknown")),
+                            "error": f"Serialization failed: {str(e)}",
+                        }
+                    )
         return {
             "session_id": self.session_id,
             "user_id": self.user_id,
             "workflow_id": self.workflow_id,
             "workflow_name": self.workflow_name,
-            "runs": [run.to_dict() for run in self.runs] if self.runs else None,
+            "runs": runs_data,
             "session_data": self.session_data,
             "workflow_data": self.workflow_data,
             "extra_data": self.extra_data,

agno/storage/singlestore.py

@@ -120,6 +120,7 @@ class SingleStoreStorage(Storage):
             specific_columns = [
                 Column("workflow_id", mysql.TEXT),
                 Column("workflow_data", mysql.JSON),
+                Column("runs", mysql.JSON),
             ]
 
         # Create table with all columns
@@ -518,7 +519,9 @@ class SingleStoreStorage(Storage):
                             "workflow_id": session.workflow_id,  # type: ignore
                             "user_id": session.user_id,
                             "workflow_name": session.workflow_name,  # type: ignore
-                            "runs": session_dict.get("runs"),
+                            "runs": json.dumps(session_dict.get("runs"), ensure_ascii=False)
+                            if session_dict.get("runs")
+                            else None,
                             "workflow_data": json.dumps(session.workflow_data, ensure_ascii=False)  # type: ignore
                             if session.workflow_data is not None  # type: ignore
                             else None,

agno/storage/sqlite.py

@@ -137,7 +137,6 @@ class SqliteStorage(Storage):
                 Column("workflow_id", String, index=True),
                 Column("workflow_name", String, index=True),
                 Column("workflow_data", sqlite.JSON),
-                # Store WorkflowRunResponse objects
                 Column("runs", sqlite.JSON),
             ]
 

agno/team/team.py

@@ -688,15 +688,6 @@ class Team:
             if self.workflow_session_state is not None:
                 self.workflow_session_state["current_user_id"] = user_id
 
-    def _reset_session_state(self) -> None:
-        """Reset the session state for the agent."""
-        if self.team_session_state is not None:
-            self.team_session_state.pop("current_session_id", None)
-            self.team_session_state.pop("current_user_id", None)
-        if self.session_state is not None:
-            self.session_state.pop("current_session_id", None)
-            self.session_state.pop("current_user_id", None)
-
     def _initialize_session(
         self,
         session_id: Optional[str] = None,
@@ -974,8 +965,6 @@ class Team:
                         from_run_response=run_response,
                         session_id=session_id,
                     )
-            finally:
-                self._reset_session_state()
 
         # If we get here, all retries failed
         if last_exception is not None:
@@ -1363,8 +1352,6 @@ class Team:
                         from_run_response=run_response,
                         session_id=session_id,
                     )
-            finally:
-                self._reset_session_state()
 
         # If we get here, all retries failed
         if last_exception is not None:
@@ -5892,7 +5879,7 @@ class Team:
 
             Args:
                 task_description (str): The task description to send to the member agents.
-                expected_output (str): The expected output from the member agents.
+                expected_output (str, optional): The expected output from the member agents.
 
             Returns:
                 str: The responses from the member agents.
@@ -5906,14 +5893,18 @@ class Team:
                 session_id, images, videos, audio
             )
 
-            # 3. Create the member agent task
-            member_agent_task = self._format_member_agent_task(
-                task_description, expected_output, team_context_str, team_member_interactions_str
-            )
-
+            # 3. Run members
             for member_agent_index, member_agent in enumerate(self.members):
                 self._initialize_member(member_agent, session_id=session_id)
 
+                # Don't override the expected output of a member agent
+                if member_agent.expected_output is not None:
+                    expected_output = None
+
+                member_agent_task = self._format_member_agent_task(
+                    task_description, expected_output, team_context_str, team_member_interactions_str
+                )
+
                 if stream:
                     member_agent_run_response_stream = member_agent.run(
                         member_agent_task,
@@ -6021,11 +6012,6 @@ class Team:
                 session_id, images, videos, audio
             )
 
-            # 3. Create the member agent task
-            member_agent_task = self._format_member_agent_task(
-                task_description, expected_output, team_context_str, team_member_interactions_str
-            )
-
             # Create tasks for all member agents
             tasks = []
             for member_agent_index, member_agent in enumerate(self.members):
@@ -6034,6 +6020,14 @@ class Team:
                 current_index = member_agent_index  # Create a reference to the current index
                 self._initialize_member(current_agent, session_id=session_id)
 
+                # Don't override the expected output of a member agent
+                if current_agent.expected_output is not None:
+                    expected_output = None
+
+                member_agent_task = self._format_member_agent_task(
+                    task_description, expected_output, team_context_str, team_member_interactions_str
+                )
+
                 async def run_member_agent(agent=current_agent, idx=current_index) -> str:
                     response = await agent.arun(
                         member_agent_task,
@@ -6185,7 +6179,7 @@ class Team:
             Args:
                 member_id (str): The ID of the member to transfer the task to. Use only the ID of the member, not the ID of the team followed by the ID of the member.
                 task_description (str): A clear and concise description of the task the member should achieve.
-                expected_output (str): The expected output from the member (optional).
+                expected_output (str, optional): The expected output from the member (optional).
             Returns:
                 str: The result of the delegated task.
             """
@@ -6204,6 +6198,9 @@ class Team:
             )
 
             # 3. Create the member agent task
+            # Don't override the expected output of a member agent
+            if member_agent.expected_output is not None:
+                expected_output = None
             member_agent_task = self._format_member_agent_task(
                 task_description, expected_output, team_context_str, team_member_interactions_str
             )
@@ -6324,7 +6321,7 @@ class Team:
             Args:
                 member_id (str): The ID of the member to transfer the task to. Use only the ID of the member, not the ID of the team followed by the ID of the member.
                 task_description (str): A clear and concise description of the task the member should achieve.
-                expected_output (str): The expected output from the member (optional).
+                expected_output (str, optional): The expected output from the member (optional).
             Returns:
                 str: The result of the delegated task.
             """
@@ -6344,6 +6341,9 @@ class Team:
             )
 
             # 3. Create the member agent task
+            # Don't override the expected output of a member agent
+            if member_agent.expected_output is not None:
+                expected_output = None
             member_agent_task = self._format_member_agent_task(
                 task_description, expected_output, team_context_str, team_member_interactions_str
             )
@@ -6555,7 +6555,7 @@ class Team:
             """Use this function to forward the request to the selected team member.
             Args:
                 member_id (str): The ID of the member to transfer the task to. Use only the ID of the member, not the ID of the team followed by the ID of the member.
-                expected_output (str): The expected output from the member (optional).
+                expected_output (str, optional): The expected output from the member (optional).
             Returns:
                 str: The result of the delegated task.
             """
@@ -6584,7 +6584,8 @@ class Team:
             # If found in subteam, include the path in the task description
             member_agent_task = message.get_content_string()
 
-            if expected_output:
+            # Don't override the expected output of a member agent
+            if member_agent.expected_output is None and expected_output:
                 member_agent_task += f"\n\n<expected_output>\n{expected_output}\n</expected_output>"
 
             # Handle enable_agentic_knowledge_filters
@@ -6694,7 +6695,7 @@ class Team:
 
             Args:
                 member_id (str): The ID of the member to transfer the task to. Use only the ID of the member, not the ID of the team followed by the ID of the member.
-                expected_output (str): The expected output from the member (optional).
+                expected_output (str, optional): The expected output from the member (optional).
             Returns:
                 str: The result of the delegated task.
             """
@@ -6719,7 +6720,8 @@ class Team:
             # If found in subteam, include the path in the task description
             member_agent_task = message.get_content_string()
 
-            if expected_output:
+            # Don't override the expected output of a member agent
+            if member_agent.expected_output is None and expected_output:
                 member_agent_task += f"\n\n<expected_output>\n{expected_output}\n</expected_output>"
 
             # Handle enable_agentic_knowledge_filters
@@ -7370,6 +7372,9 @@ class Team:
         """Return a list of references from the knowledge base"""
         from agno.document import Document
 
+        if num_documents is None and self.knowledge is not None:
+            num_documents = self.knowledge.num_documents
+
         # Validate the filters against known valid filter keys
         if self.knowledge is not None:
             valid_filters, invalid_keys = self.knowledge.validate_filters(filters)  # type: ignore
@@ -7404,9 +7409,6 @@ class Team:
             if self.knowledge is None or self.knowledge.vector_db is None:
                 return None
 
-            if num_documents is None:
-                num_documents = self.knowledge.num_documents
-
             log_debug(f"Searching knowledge base with filters: {filters}")
             relevant_docs: List[Document] = self.knowledge.search(
                 query=query, num_documents=num_documents, filters=filters
@@ -7427,6 +7429,9 @@ class Team:
         """Get relevant documents from knowledge base asynchronously."""
         from agno.document import Document
 
+        if num_documents is None and self.knowledge is not None:
+            num_documents = self.knowledge.num_documents
+
         # Validate the filters against known valid filter keys
         if self.knowledge is not None:
             valid_filters, invalid_keys = self.knowledge.validate_filters(filters)  # type: ignore
@@ -7463,9 +7468,6 @@ class Team:
             if self.knowledge is None or self.knowledge.vector_db is None:
                 return None
 
-            if num_documents is None:
-                num_documents = self.knowledge.num_documents
-
             log_debug(f"Searching knowledge base with filters: {filters}")
             relevant_docs: List[Document] = await self.knowledge.async_search(
                 query=query, num_documents=num_documents, filters=filters