agno 2.2.13__py3-none-any.whl → 2.4.3__py3-none-any.whl

agno/tools/workflow.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, Optional
 
 from pydantic import BaseModel, Field
 
+from agno.run import RunContext
 from agno.tools import Toolkit
 from agno.utils.log import log_debug, log_error
 from agno.workflow.workflow import Workflow, WorkflowRunOutput
@@ -65,7 +66,7 @@ class WorkflowTools(Toolkit):
             else:
                 self.register(self.analyze, name="analyze")
 
-    def think(self, session_state: Dict[str, Any], thought: str) -> str:
+    def think(self, run_context: RunContext, thought: str) -> str:
         """Use this tool as a scratchpad to reason about the workflow execution, refine your approach, brainstorm workflow inputs, or revise your plan.
         Call `Think` whenever you need to figure out what to do next, analyze the user's requirements, plan workflow inputs, or decide on execution strategy.
         You should use this tool as frequently as needed.
@@ -76,14 +77,14 @@ class WorkflowTools(Toolkit):
             log_debug(f"Workflow Thought: {thought}")
 
             # Add the thought to the session state
-            if session_state is None:
-                session_state = {}
-            if "workflow_thoughts" not in session_state:
-                session_state["workflow_thoughts"] = []
-            session_state["workflow_thoughts"].append(thought)
+            if run_context.session_state is None:
+                run_context.session_state = {}
+            if "workflow_thoughts" not in run_context.session_state:
+                run_context.session_state["workflow_thoughts"] = []
+            run_context.session_state["workflow_thoughts"].append(thought)
 
             # Return the full log of thoughts and the new thought
-            thoughts = "\n".join([f"- {t}" for t in session_state["workflow_thoughts"]])
+            thoughts = "\n".join([f"- {t}" for t in run_context.session_state["workflow_thoughts"]])
             formatted_thoughts = dedent(
                 f"""Workflow Thoughts:
                 {thoughts}
@@ -94,7 +95,7 @@ class WorkflowTools(Toolkit):
             log_error(f"Error recording workflow thought: {e}")
             return f"Error recording workflow thought: {e}"
 
-    async def async_think(self, session_state: Dict[str, Any], thought: str) -> str:
+    async def async_think(self, run_context: RunContext, thought: str) -> str:
         """Use this tool as a scratchpad to reason about the workflow execution, refine your approach, brainstorm workflow inputs, or revise your plan.
         Call `Think` whenever you need to figure out what to do next, analyze the user's requirements, plan workflow inputs, or decide on execution strategy.
         You should use this tool as frequently as needed.
@@ -105,14 +106,14 @@ class WorkflowTools(Toolkit):
             log_debug(f"Workflow Thought: {thought}")
 
             # Add the thought to the session state
-            if session_state is None:
-                session_state = {}
-            if "workflow_thoughts" not in session_state:
-                session_state["workflow_thoughts"] = []
-            session_state["workflow_thoughts"].append(thought)
+            if run_context.session_state is None:
+                run_context.session_state = {}
+            if "workflow_thoughts" not in run_context.session_state:
+                run_context.session_state["workflow_thoughts"] = []
+            run_context.session_state["workflow_thoughts"].append(thought)
 
             # Return the full log of thoughts and the new thought
-            thoughts = "\n".join([f"- {t}" for t in session_state["workflow_thoughts"]])
+            thoughts = "\n".join([f"- {t}" for t in run_context.session_state["workflow_thoughts"]])
             formatted_thoughts = dedent(
                 f"""Workflow Thoughts:
                 {thoughts}
@@ -125,33 +126,37 @@ class WorkflowTools(Toolkit):
 
     def run_workflow(
         self,
-        session_state: Dict[str, Any],
+        run_context: RunContext,
         input: RunWorkflowInput,
     ) -> str:
         """Use this tool to execute the workflow with the specified inputs and parameters.
         After thinking through the requirements, use this tool to run the workflow with appropriate inputs.
+
         Args:
-            input_data: The input data for the workflow.
+            input: The input data for the workflow.
         """
+        if isinstance(input, dict):
+            input = RunWorkflowInput.model_validate(input)
+
         try:
             log_debug(f"Running workflow with input: {input.input_data}")
 
-            user_id = session_state.get("current_user_id")
-            session_id = session_state.get("current_session_id")
+            if run_context.session_state is None:
+                run_context.session_state = {}
 
             # Execute the workflow
             result: WorkflowRunOutput = self.workflow.run(
                 input=input.input_data,
-                user_id=user_id,
-                session_id=session_id,
-                session_state=session_state,
+                user_id=run_context.user_id,
+                session_id=run_context.session_id,
+                session_state=run_context.session_state,
                 additional_data=input.additional_data,
             )
 
-            if "workflow_results" not in session_state:
-                session_state["workflow_results"] = []
+            if "workflow_results" not in run_context.session_state:
+                run_context.session_state["workflow_results"] = []
 
-            session_state["workflow_results"].append(result.to_dict())
+            run_context.session_state["workflow_results"].append(result.to_dict())
 
             return json.dumps(result.to_dict(), indent=2)
 
@@ -161,7 +166,7 @@ class WorkflowTools(Toolkit):
 
     async def async_run_workflow(
         self,
-        session_state: Dict[str, Any],
+        run_context: RunContext,
         input: RunWorkflowInput,
     ) -> str:
         """Use this tool to execute the workflow with the specified inputs and parameters.
@@ -170,25 +175,28 @@ class WorkflowTools(Toolkit):
             input_data: The input data for the workflow (use a `str` for a simple input)
             additional_data: The additional data for the workflow. This is a dictionary of key-value pairs that will be passed to the workflow. E.g. {"topic": "food", "style": "Humour"}
         """
+        if isinstance(input, dict):
+            input = RunWorkflowInput.model_validate(input)
+
         try:
             log_debug(f"Running workflow with input: {input.input_data}")
 
-            user_id = session_state.get("current_user_id")
-            session_id = session_state.get("current_session_id")
+            if run_context.session_state is None:
+                run_context.session_state = {}
 
             # Execute the workflow
             result: WorkflowRunOutput = await self.workflow.arun(
                 input=input.input_data,
-                user_id=user_id,
-                session_id=session_id,
-                session_state=session_state,
+                user_id=run_context.user_id,
+                session_id=run_context.session_id,
+                session_state=run_context.session_state,
                 additional_data=input.additional_data,
             )
 
-            if "workflow_results" not in session_state:
-                session_state["workflow_results"] = []
+            if "workflow_results" not in run_context.session_state:
+                run_context.session_state["workflow_results"] = []
 
-            session_state["workflow_results"].append(result.to_dict())
+            run_context.session_state["workflow_results"].append(result.to_dict())
 
             return json.dumps(result.to_dict(), indent=2)
 
@@ -196,7 +204,7 @@ class WorkflowTools(Toolkit):
             log_error(f"Error running workflow: {e}")
             return f"Error running workflow: {e}"
 
-    def analyze(self, session_state: Dict[str, Any], analysis: str) -> str:
+    def analyze(self, run_context: RunContext, analysis: str) -> str:
         """Use this tool to evaluate whether the workflow execution results are correct and sufficient.
         If not, go back to "Think" or "Run" with refined inputs or parameters.
         Args:
@@ -206,14 +214,14 @@ class WorkflowTools(Toolkit):
             log_debug(f"Workflow Analysis: {analysis}")
 
             # Add the analysis to the session state
-            if session_state is None:
-                session_state = {}
-            if "workflow_analysis" not in session_state:
-                session_state["workflow_analysis"] = []
-            session_state["workflow_analysis"].append(analysis)
+            if run_context.session_state is None:
+                run_context.session_state = {}
+            if "workflow_analysis" not in run_context.session_state:
+                run_context.session_state["workflow_analysis"] = []
+            run_context.session_state["workflow_analysis"].append(analysis)
 
             # Return the full log of analysis and the new analysis
-            analysis_log = "\n".join([f"- {a}" for a in session_state["workflow_analysis"]])
+            analysis_log = "\n".join([f"- {a}" for a in run_context.session_state["workflow_analysis"]])
             formatted_analysis = dedent(
                 f"""Workflow Analysis:
                 {analysis_log}
@@ -224,7 +232,7 @@ class WorkflowTools(Toolkit):
             log_error(f"Error recording workflow analysis: {e}")
             return f"Error recording workflow analysis: {e}"
 
-    async def async_analyze(self, session_state: Dict[str, Any], analysis: str) -> str:
+    async def async_analyze(self, run_context: RunContext, analysis: str) -> str:
         """Use this tool to evaluate whether the workflow execution results are correct and sufficient.
         If not, go back to "Think" or "Run" with refined inputs or parameters.
         Args:
@@ -234,14 +242,14 @@ class WorkflowTools(Toolkit):
             log_debug(f"Workflow Analysis: {analysis}")
 
             # Add the analysis to the session state
-            if session_state is None:
-                session_state = {}
-            if "workflow_analysis" not in session_state:
-                session_state["workflow_analysis"] = []
-            session_state["workflow_analysis"].append(analysis)
+            if run_context.session_state is None:
+                run_context.session_state = {}
+            if "workflow_analysis" not in run_context.session_state:
+                run_context.session_state["workflow_analysis"] = []
+            run_context.session_state["workflow_analysis"].append(analysis)
 
             # Return the full log of analysis and the new analysis
-            analysis_log = "\n".join([f"- {a}" for a in session_state["workflow_analysis"]])
+            analysis_log = "\n".join([f"- {a}" for a in run_context.session_state["workflow_analysis"]])
             formatted_analysis = dedent(
                 f"""Workflow Analysis:
                 {analysis_log}
@@ -255,7 +263,7 @@ class WorkflowTools(Toolkit):
     DEFAULT_INSTRUCTIONS = dedent("""\
         You have access to the Think, Run Workflow, and Analyze tools that will help you execute workflows and analyze their results. Use these tools as frequently as needed to successfully complete workflow-based tasks.
         ## How to use the Think, Run Workflow, and Analyze tools:
-        
+
         1. **Think**
         - Purpose: A scratchpad for planning workflow execution, brainstorming inputs, and refining your approach. You never reveal your "Think" content to the user.
         - Usage: Call `think` whenever you need to figure out what workflow inputs to use, analyze requirements, or decide on execution strategy before (or after) you run the workflow.

agno/tools/yfinance.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, List
+from typing import Any, List, Optional
 
 from agno.tools import Toolkit
 from agno.utils.log import log_debug
@@ -18,6 +18,7 @@ class YFinanceTools(Toolkit):
 
     def __init__(
         self,
+        session: Optional[Any] = None,
         **kwargs,
     ):
         tools: List[Any] = [
@@ -31,7 +32,7 @@ class YFinanceTools(Toolkit):
             self.get_technical_indicators,
             self.get_historical_stock_prices,
         ]
-
+        self.session = session
         super().__init__(name="yfinance_tools", tools=tools, **kwargs)
 
     def get_current_stock_price(self, symbol: str) -> str:
@@ -46,7 +47,7 @@ class YFinanceTools(Toolkit):
         """
         try:
             log_debug(f"Fetching current price for {symbol}")
-            stock = yf.Ticker(symbol)
+            stock = yf.Ticker(symbol, session=self.session)
             # Use "regularMarketPrice" for regular market hours, or "currentPrice" for pre/post market
             current_price = stock.info.get("regularMarketPrice", stock.info.get("currentPrice"))
             return f"{current_price:.4f}" if current_price else f"Could not fetch current price for {symbol}"
@@ -63,7 +64,7 @@ class YFinanceTools(Toolkit):
             str: JSON containing company profile and overview.
         """
         try:
-            company_info_full = yf.Ticker(symbol).info
+            company_info_full = yf.Ticker(symbol, session=self.session).info
             if company_info_full is None:
                 return f"Could not fetch company info for {symbol}"
 
@@ -120,7 +121,7 @@ class YFinanceTools(Toolkit):
         """
         try:
             log_debug(f"Fetching historical prices for {symbol}")
-            stock = yf.Ticker(symbol)
+            stock = yf.Ticker(symbol, session=self.session)
             historical_price = stock.history(period=period, interval=interval)
             return historical_price.to_json(orient="index")
         except Exception as e:
@@ -150,7 +151,7 @@ class YFinanceTools(Toolkit):
         """
         try:
             log_debug(f"Fetching fundamentals for {symbol}")
-            stock = yf.Ticker(symbol)
+            stock = yf.Ticker(symbol, session=self.session)
             info = stock.info
             fundamentals = {
                 "symbol": symbol,
@@ -181,7 +182,7 @@ class YFinanceTools(Toolkit):
         """
         try:
             log_debug(f"Fetching income statements for {symbol}")
-            stock = yf.Ticker(symbol)
+            stock = yf.Ticker(symbol, session=self.session)
             financials = stock.financials
             return financials.to_json(orient="index")
         except Exception as e:
@@ -198,7 +199,7 @@ class YFinanceTools(Toolkit):
         """
         try:
             log_debug(f"Fetching key financial ratios for {symbol}")
-            stock = yf.Ticker(symbol)
+            stock = yf.Ticker(symbol, session=self.session)
             key_ratios = stock.info
             return json.dumps(key_ratios, indent=2)
         except Exception as e:
@@ -215,7 +216,7 @@ class YFinanceTools(Toolkit):
         """
         try:
             log_debug(f"Fetching analyst recommendations for {symbol}")
-            stock = yf.Ticker(symbol)
+            stock = yf.Ticker(symbol, session=self.session)
             recommendations = stock.recommendations
             return recommendations.to_json(orient="index")
         except Exception as e:
@@ -233,7 +234,7 @@ class YFinanceTools(Toolkit):
         """
         try:
             log_debug(f"Fetching company news for {symbol}")
-            news = yf.Ticker(symbol).news
+            news = yf.Ticker(symbol, session=self.session).news
             return json.dumps(news[:num_stories], indent=2)
         except Exception as e:
             return f"Error fetching company news for {symbol}: {e}"
@@ -251,7 +252,7 @@ class YFinanceTools(Toolkit):
         """
         try:
             log_debug(f"Fetching technical indicators for {symbol}")
-            indicators = yf.Ticker(symbol).history(period=period)
+            indicators = yf.Ticker(symbol, session=self.session).history(period=period)
             return indicators.to_json(orient="index")
         except Exception as e:
             return f"Error fetching technical indicators for {symbol}: {e}"

agno/tracing/__init__.py

@@ -0,0 +1,12 @@
+"""
+Agno Tracing Module
+
+This module provides OpenTelemetry-based tracing capabilities for Agno agents.
+It uses the openinference-instrumentation-agno package for automatic instrumentation
+and provides a custom DatabaseSpanExporter to store traces in the Agno database.
+"""
+
+from agno.tracing.exporter import DatabaseSpanExporter
+from agno.tracing.setup import setup_tracing
+
+__all__ = ["DatabaseSpanExporter", "setup_tracing"]

agno/tracing/exporter.py

@@ -0,0 +1,161 @@
+"""
+Custom OpenTelemetry SpanExporter that writes traces to Agno database.
+"""
+
+import asyncio
+from collections import defaultdict
+from typing import Dict, List, Sequence, Union
+
+from opentelemetry.sdk.trace import ReadableSpan  # type: ignore
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult  # type: ignore
+
+from agno.db.base import AsyncBaseDb, BaseDb
+from agno.remote.base import RemoteDb
+from agno.tracing.schemas import Span, create_trace_from_spans
+from agno.utils.log import logger
+
+
+class DatabaseSpanExporter(SpanExporter):
+    """Custom OpenTelemetry SpanExporter that writes to Agno database"""
+
+    def __init__(self, db: Union[BaseDb, AsyncBaseDb, RemoteDb]):
+        """
+        Initialize the DatabaseSpanExporter.
+
+        Args:
+            db: Database instance (sync or async) to store traces
+        """
+        self.db = db
+        self._shutdown = False
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """
+        Export spans to the database.
+
+        This method:
+        1. Converts OpenTelemetry spans to Span objects
+        2. Groups spans by trace_id
+        3. Creates Trace records (one per trace_id)
+        4. Creates Span records (multiple per trace_id)
+
+        Args:
+            spans: Sequence of OpenTelemetry ReadableSpan objects
+
+        Returns:
+            SpanExportResult indicating success or failure
+        """
+        if self._shutdown:
+            logger.warning("DatabaseSpanExporter is shutdown, cannot export spans")
+            return SpanExportResult.FAILURE
+
+        if not spans:
+            return SpanExportResult.SUCCESS
+
+        try:
+            # Convert OpenTelemetry spans to Span objects
+            converted_spans: List[Span] = []
+            for span in spans:
+                try:
+                    converted_span = Span.from_otel_span(span)
+                    converted_spans.append(converted_span)
+                except Exception as e:
+                    logger.error(f"Failed to convert span {span.name}: {e}")
+                    # Continue processing other spans
+                    continue
+
+            if not converted_spans:
+                return SpanExportResult.SUCCESS
+
+            # Group spans by trace_id
+            spans_by_trace: Dict[str, List[Span]] = defaultdict(list)
+            for converted_span in converted_spans:
+                spans_by_trace[converted_span.trace_id].append(converted_span)
+
+            # Handle async DB
+            if isinstance(self.db, RemoteDb):
+                # Skipping remote database because it handles its own tracing
+                pass
+            elif isinstance(self.db, AsyncBaseDb):
+                self._export_async(spans_by_trace)
+            else:
+                # Synchronous database
+                self._export_sync(spans_by_trace)
+
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            logger.error(f"Failed to export spans to database: {e}", exc_info=True)
+            return SpanExportResult.FAILURE
+
+    def _export_sync(self, spans_by_trace: Dict[str, List[Span]]) -> None:
+        """Export traces and spans to synchronous database"""
+        try:
+            # Create trace and span records for each trace
+            for trace_id, spans in spans_by_trace.items():
+                # Create trace record (aggregate of all spans)
+                trace = create_trace_from_spans(spans)
+                if trace:
+                    self.db.upsert_trace(trace)  # type: ignore
+
+                # Create span records
+                self.db.create_spans(spans)  # type: ignore
+
+        except Exception as e:
+            logger.error(f"Failed to export sync traces: {e}", exc_info=True)
+            raise
+
+    def _export_async(self, spans_by_trace: Dict[str, List[Span]]) -> None:
+        """Handle async database export"""
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # We're in an async context, schedule the coroutine
+                asyncio.create_task(self._do_async_export(spans_by_trace))
+            else:
+                # No running loop, run in new loop
+                loop.run_until_complete(self._do_async_export(spans_by_trace))
+        except RuntimeError:
+            # No event loop, create new one
+            try:
+                asyncio.run(self._do_async_export(spans_by_trace))
+            except Exception as e:
+                logger.error(f"Failed to export async traces: {e}", exc_info=True)
+
+    async def _do_async_export(self, spans_by_trace: Dict[str, List[Span]]) -> None:
+        """Actually perform the async export"""
+        try:
+            # Create trace and span records for each trace
+            for trace_id, spans in spans_by_trace.items():
+                # Create trace record (aggregate of all spans)
+                trace = create_trace_from_spans(spans)
+                if trace:
+                    create_trace_result = self.db.upsert_trace(trace)  # type: ignore
+                    if create_trace_result is not None:
+                        await create_trace_result
+
+                # Create span records
+                create_spans_result = self.db.create_spans(spans)  # type: ignore
+                if create_spans_result is not None:
+                    await create_spans_result
+
+        except Exception as e:
+            logger.error(f"Failed to do async export: {e}", exc_info=True)
+            raise
+
+    def shutdown(self) -> None:
+        """Shutdown the exporter"""
+        self._shutdown = True
+        logger.debug("DatabaseSpanExporter shutdown")
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """
+        Force flush any pending spans.
+
+        Since we write immediately to the database, this is a no-op.
+
+        Args:
+            timeout_millis: Timeout in milliseconds
+
+        Returns:
+            True if flush was successful
+        """
+        return True