quraite 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

quraite/__init__.py

@@ -1,3 +1,7 @@
 """Quraite Python SDK"""
 
+from quraite.serve import create_app, run_agent, setup_tunnel
+
 __version__ = "0.6.0"
+
+__all__ = ["create_app", "run_agent", "setup_tunnel"]

quraite/adapters/agno_adapter.py

@@ -1,135 +1,93 @@
-from typing import List, Union
+import uuid
 
 from agno.agent import Agent
 from agno.team import Team
+from openinference.instrumentation.agno.utils import _AGNO_PARENT_NODE_CONTEXT_KEY
+from opentelemetry import context as context_api
 from opentelemetry.trace import TracerProvider
 
 from quraite.adapters.base import BaseAdapter
+from quraite.constants.framework import Framework
 from quraite.logger import get_logger
-from quraite.schema.message import AgentMessage
-from quraite.schema.response import AgentInvocationResponse
-from quraite.tracing.constants import Framework
+from quraite.schema.invoke import InvokeInput, InvokeOutput
 from quraite.tracing.trace import AgentSpan, AgentTrace
 
 logger = get_logger(__name__)
 
 
 class AgnoAdapter(BaseAdapter):
-    """
-    Agno adapter wrapper that converts any Agno agent or team
-    to a standardized callable interface (invoke) and converts the output to List[AgentMessage].
-
-    This class wraps any Agno Agent or Team and provides:
-    - Asynchronous invocation via ainvoke()
-    - Automatic conversion to List[AgentMessage] format
-    - Access to message history and traces
-    - Support for both single agents and multi-agent teams
-    """
+    """Agno adapter for Agent and Team (requires tracing)."""
 
     def __init__(
         self,
-        agent: Union[Agent, Team],
+        agent: Agent | Team,
+        *,
+        tracer_provider: TracerProvider | None = None,
         agent_name: str = "Agno Agent",
-        tracer_provider: TracerProvider = None,
     ):
         """
-        Initialize with a pre-configured Agno agent or team
+        Initialize Agno adapter.
 
         Args:
-            agent: An Agno Agent or Team instance
-            agent_name: Name of the agent for trajectory metadata
-            tracer_provider: TracerProvider for tracing (required)
+            agent: Agno Agent or Team instance
+            tracer_provider: TracerProvider from setup_tracing() (required)
+            agent_name: Agent name for metadata
         """
-        logger.debug("Initializing AgnoAdapter with agent_name=%s", agent_name)
-        self._init_tracing(tracer_provider, required=True)
+        if tracer_provider is None:
+            raise ValueError(
+                "Agno adapter requires tracing. Use setup_tracing([Framework.AGNO]) first."
+            )
 
-        self.agent: Union[Agent, Team] = agent
+        self.agent: Agent | Team = agent
         self.agent_name = agent_name
-        logger.info(
-            "AgnoAdapter initialized successfully (tracing_enabled=%s)",
-            bool(tracer_provider),
-        )
-
-    def _prepare_input(self, input: List[AgentMessage]) -> str:
-        """
-        Prepare input for Agno agent from List[AgentMessage].
-
-        Args:
-            input: List[AgentMessage] containing user_message
-
-        Returns:
-            str: User message text
-        """
-        logger.debug("Preparing input from %d messages", len(input))
-        if not input or input[-1].role != "user":
-            logger.error("Invalid input: no user message found")
-            raise ValueError("No user message found in the input")
-
-        last_user_message = input[-1]
-        if not last_user_message.content:
-            logger.error("User message has no content")
-            raise ValueError("User message has no content")
-
-        text_content = None
-        for content_item in last_user_message.content:
-            if content_item.type == "text" and content_item.text:
-                text_content = content_item.text
-                break
-
-        if not text_content:
-            logger.error("No text content found in user message")
-            raise ValueError("No text content found in user message")
+        self._init_tracer(tracer_provider)
+        logger.info("AgnoAdapter initialized")
 
-        logger.debug("Prepared input with text_content length=%d", len(text_content))
-        return text_content
-
-    async def ainvoke(
-        self,
-        input: List[AgentMessage],
-        session_id: Union[str, None] = None,
-    ) -> AgentInvocationResponse:
-        """
-        Asynchronous invocation method - invokes the Agno agent/team with tracing
-
-        Args:
-            input: List[AgentMessage] containing user_message
-            session_id: Optional conversation ID for maintaining context
-
-        Returns:
-            AgentInvocationResponse - response containing agent trace, trajectory, and final response.
-        """
-        logger.info(
-            "ainvoke called (session_id=%s, input_messages=%d)", session_id, len(input)
-        )
-        agent_input = self._prepare_input(input)
-        session_id = session_id or "default"
+    async def ainvoke(self, input: InvokeInput) -> InvokeOutput:
+        """Invoke Agno agent and return response with trace."""
+        logger.info("ainvoke called (session_id=%s)", input.session_id)
+        session_id = input.session_id or str(uuid.uuid4())
 
         try:
-            logger.debug("Invoking Agno agent with tracing (session_id=%s)", session_id)
-            return await self._ainvoke_with_tracing(agent_input, session_id)
-
-        except ValueError:
-            logger.exception("Invalid input during ainvoke")
-            raise
+            return await self._ainvoke_with_tracing(
+                input.user_message_str(), session_id
+            )
         except Exception as e:
-            logger.exception("Unexpected error invoking Agno agent")
+            logger.exception("Error invoking Agno agent")
             raise RuntimeError(f"Error invoking Agno agent: {e}") from e
 
     async def _ainvoke_with_tracing(
         self,
-        agent_input: str,
+        user_message: str,
         session_id: str,
-    ) -> AgentInvocationResponse:
+    ) -> InvokeOutput:
         """Execute ainvoke with tracing enabled."""
         with self.tracer.start_as_current_span("agno_invocation") as span:
+            # Workaround: openinference-instrumentation-agno>=0.1.25 forces a new root
+            # trace for Team instances unless _AGNO_PARENT_NODE_CONTEXT_KEY is set.
+            # This ensures child spans inherit our trace_id.
+            # See: https://github.com/Arize-ai/openinference/pull/2533
+            #
+            # Context API notes:
+            # - set_value() creates a new Context with the key-value pair (doesn't activate it)
+            # - attach() makes the context active and returns a token (restore point)
+            # - detach(token) restores the previous context (cleanup)
+            ctx = context_api.set_value(
+                _AGNO_PARENT_NODE_CONTEXT_KEY,
+                format(span.get_span_context().span_id, '016x')
+            )
+            token = context_api.attach(ctx)
             trace_id = span.get_span_context().trace_id
+
             logger.debug(
                 "Starting traced invocation (session_id=%s) with trace_id=%s",
                 session_id,
                 trace_id,
             )
-            # Run the agent/team
-            await self.agent.arun(agent_input, session_id=session_id)
+            try:
+                await self.agent.arun(user_message, session_id=session_id)
+            finally:
+                context_api.detach(token)
 
         # Get trace spans
         trace_readable_spans = self.quraite_span_exporter.get_spans_by_trace_id(
@@ -151,7 +109,7 @@ class AgnoAdapter(BaseAdapter):
         else:
             logger.warning("No spans found for trace_id=%s", trace_id)
 
-        return AgentInvocationResponse(
+        return InvokeOutput(
             agent_trace=agent_trace,
             agent_trajectory=agent_trace.to_agent_trajectory(framework=Framework.AGNO),
         )

quraite/adapters/base.py

@@ -1,49 +1,28 @@
 from abc import ABC, abstractmethod
-from typing import Any, List, Optional, Union
 
+from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider
 from opentelemetry.trace import Tracer, TracerProvider
 
-from quraite.schema.message import AgentMessage, AssistantMessage, MessageContentText
-from quraite.schema.response import AgentInvocationResponse
+from quraite.logger import get_logger
+from quraite.schema.invoke import InvokeInput, InvokeOutput
+from quraite.schema.message import AssistantMessage, MessageContentText
 from quraite.tracing.constants import QURAITE_TRACER_NAME
-from quraite.tracing.span_exporter import QuraiteInMemorySpanExporter
-from quraite.tracing.span_processor import QuraiteSimpleSpanProcessor
+from quraite.tracing.span_exporter import QuraiteSpanExporter
+from quraite.tracing.span_processor import QuraiteSpanProcessor
 
+logger = get_logger(__name__)
 
-class BaseAdapter(ABC):
-    """
-    Abstract base class for all adapter providers.
-
-    Subclasses must implement the asynchronous adapter method,
-    which take List[Message] and return agent Message outputs.
-
-    Methods:
-        ainvoke: Asynchronously invoke an agent call for the provided input.
-    """
 
-    tracer_provider: Optional[TracerProvider] = None
-    tracer: Optional[Tracer] = None
-    quraite_span_exporter: Optional[QuraiteInMemorySpanExporter] = None
+class BaseAdapter(ABC):
+    """Base adapter for all framework adapters."""
 
-    def _init_tracing(
-        self,
-        tracer_provider: Optional[TracerProvider],
-        required: bool = False,
-    ) -> None:
-        """
-        Initialize tracing components from a TracerProvider.
+    tracer_provider: TracerProvider | None = None
+    tracer: Tracer | None = None
+    quraite_span_exporter: QuraiteSpanExporter = QuraiteSpanExporter()
 
-        Args:
-            tracer_provider: TracerProvider for tracing
-            span_exporter: SpanExporter for exporting spans
-            required: If True, raises ValueError when tracer_provider is None
-        """
+    def _init_tracer(self, tracer_provider: SDKTracerProvider | None) -> None:
+        """Initialize tracer components from TracerProvider."""
         if tracer_provider is None:
-            if required:
-                raise ValueError(
-                    "tracer_provider is required. "
-                    "Please provide a TracerProvider instance."
-                )
             return
 
         self.tracer_provider = tracer_provider
@@ -54,7 +33,7 @@ class BaseAdapter(ABC):
             (
                 processor.span_exporter
                 for processor in tracer_provider._active_span_processor._span_processors
-                if isinstance(processor, QuraiteSimpleSpanProcessor)
+                if isinstance(processor, QuraiteSpanProcessor)
             ),
             None,
         )
@@ -62,7 +41,7 @@ class BaseAdapter(ABC):
         if quraite_span_exporter is None:
             raise ValueError(
                 "Quraite span exporter not found. "
-                "Please ensure QuraiteSimpleSpanProcessor is used in the tracer provider."
+                "Use setup_tracing() to configure tracing properly."
             )
 
         self.quraite_span_exporter = quraite_span_exporter
@@ -70,54 +49,25 @@ class BaseAdapter(ABC):
     @abstractmethod
     async def ainvoke(
         self,
-        input: List[AgentMessage],
-        session_id: Union[str, None],
-    ) -> AgentInvocationResponse:
-        """
-        Asynchronously invoke the agent with the given input.
-
-        Args:
-            input (List[AgentMessage]): List of AgentMessage objects.
-            session_id (str or None): ID for conversation thread/context.
-
-        Returns:
-            AgentInvocationResponse: Response containing agent trace, trajectory, and final response.
-
-        Raises:
-            NotImplementedError: If not implemented by subclass.
-        """
+        input: InvokeInput,
+    ) -> InvokeOutput:
+        """Invoke agent with input and return response."""
         raise NotImplementedError("Not implemented")
 
 
 class DummyAdapter(BaseAdapter):
-    """
-    A dummy implementation of BaseAdapter, used mainly for testing and scaffolding.
-    Always returns a fixed dummy response.
-
-    Methods:
-        ainvoke: Returns a static assistant message asynchronously.
-    """
+    """Dummy adapter for testing."""
 
     async def ainvoke(
         self,
-        input: Any,
-        session_id: Union[str, None],
-    ) -> AgentInvocationResponse:
-        """
-        Asynchronously returns a dummy assistant response.
-
-        Args:
-            input: Ignored.
-            session_id: Ignored.
-
-        Returns:
-            AgentInvocationResponse: Response containing agent trace, trajectory, and final response.
-        """
-
-        return AgentInvocationResponse(
+        input: InvokeInput,
+    ) -> InvokeOutput:
+        """Returns dummy response."""
+        return InvokeOutput(
             agent_trajectory=[
+                input.user_message,
                 AssistantMessage(
                     content=[MessageContentText(text="Dummy response")],
-                )
+                ),
             ]
         )

quraite/adapters/bedrock_agents_adapter.py

@@ -7,13 +7,14 @@ import asyncio
 import json
 import os
 import uuid
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 import boto3
 from botocore.exceptions import ClientError
 
 from quraite.adapters.base import BaseAdapter
 from quraite.logger import get_logger
+from quraite.schema.invoke import InvokeInput, InvokeOutput
 from quraite.schema.message import (
     AgentMessage,
     AssistantMessage,
@@ -21,7 +22,6 @@ from quraite.schema.message import (
     ToolCall,
     ToolMessage,
 )
-from quraite.schema.response import AgentInvocationResponse
 
 logger = get_logger(__name__)
 
@@ -38,12 +38,12 @@ class BedrockAgentsAdapter(BaseAdapter):
 
     def __init__(
         self,
-        aws_access_key_id: Optional[str] = None,
-        aws_secret_access_key: Optional[str] = None,
-        aws_session_token: Optional[str] = None,
-        agent_id: Optional[str] = None,
-        agent_alias_id: Optional[str] = None,
-        region_name: Optional[str] = None,
+        aws_access_key_id: str | None = None,
+        aws_secret_access_key: str | None = None,
+        aws_session_token: str | None = None,
+        agent_id: str | None = None,
+        agent_alias_id: str | None = None,
+        region_name: str | None = None,
         agent_name: str = "Bedrock Agent",
     ):
         """
@@ -92,8 +92,8 @@ class BedrockAgentsAdapter(BaseAdapter):
 
     def _convert_bedrock_traces_to_messages(
         self,
-        traces: List[Dict[str, Any]],
-    ) -> List[AgentMessage]:
+        traces: list[dict[str, Any]],
+    ) -> list[AgentMessage]:
         logger.debug("Converting %d Bedrock trace events to messages", len(traces))
         if not traces:
             return []
@@ -121,7 +121,7 @@ class BedrockAgentsAdapter(BaseAdapter):
                 try:
                     parsed_content = json.loads(raw_response_content)
                     contents = parsed_content.get("content", [])
-                except (json.JSONDecodeError, KeyError, ValueError) as e:
+                except (json.JSONDecodeError, KeyError, ValueError):
                     logger.exception("Error parsing Bedrock raw response content")
 
                 if not contents:
@@ -223,27 +223,7 @@ class BedrockAgentsAdapter(BaseAdapter):
         logger.info("Converted Bedrock traces into %d messages", len(messages))
         return messages
 
-    def _prepare_input(self, input_data: List[AgentMessage]) -> str:
-        """Extract user message from List[AgentMessage]."""
-        logger.debug("Preparing Bedrock input from %d messages", len(input_data))
-        last_user_message = input_data[-1]
-        if last_user_message.role != "user":
-            logger.error("Last message is not from user")
-            return ""
-        # Check if content list is not empty and has text
-        if not last_user_message.content:
-            logger.error("User message has no content")
-            raise ValueError("User message has no content")
-        # Find the first text content item
-        for content_item in last_user_message.content:
-            if content_item.type == "text" and content_item.text:
-                logger.debug(
-                    "Prepared Bedrock input (text_length=%d)", len(content_item.text)
-                )
-                return content_item.text
-        raise ValueError("No text content found in user message")
-
-    def _run_agent(self, session_id: str, prompt: str) -> List[Dict]:
+    def _run_agent(self, session_id: str, prompt: str) -> tuple[str, list[dict]]:
         """
         Run the Bedrock agent and collect response and traces.
 
@@ -283,7 +263,7 @@ class BedrockAgentsAdapter(BaseAdapter):
 
         except ClientError:
             logger.exception("Error invoking Bedrock agent via Bedrock runtime")
-            return "", []
+            raise
         logger.info(
             "Bedrock agent invocation succeeded (session_id=%s, trace_events=%d)",
             session_id,
@@ -291,53 +271,20 @@ class BedrockAgentsAdapter(BaseAdapter):
         )
         return agent_answer, traces
 
-    async def ainvoke(
-        self,
-        input: List[AgentMessage],
-        session_id: Union[str, None],
-    ) -> AgentInvocationResponse:
-        """Asynchronous invocation method - invokes the Bedrock agent and converts to List[AgentMessage].
-
-        Args:
-            input: List of AgentMessage objects
-            session_id: Unique session identifier
-
-        Returns:
-            AgentInvocationResponse - response containing agent trace, trajectory, and final response.
-        """
-        logger.info(
-            "Bedrock ainvoke called (session_id=%s, input_messages=%d)",
-            session_id,
-            len(input),
-        )
-        agent_input = self._prepare_input(input)
-        session_id = session_id or str(uuid.uuid4())
+    async def ainvoke(self, input: InvokeInput) -> InvokeOutput:
+        """Invoke Bedrock agent and return response."""
+        logger.info("Bedrock ainvoke called (session_id=%s)", input.session_id)
+        session_id = input.session_id or str(uuid.uuid4())
 
         try:
-            # Run the synchronous _run_agent in a thread pool to avoid blocking
             _, traces = await asyncio.to_thread(
-                self._run_agent, session_id, agent_input
-            )
-            logger.debug(
-                "Bedrock agent run returned %d trace events for session_id=%s",
-                len(traces),
-                session_id,
+                self._run_agent, session_id, input.user_message_str()
             )
-        except (ClientError, ValueError, KeyError, json.JSONDecodeError):
-            logger.exception("Error invoking Bedrock agent")
-            return AgentInvocationResponse()
-
-        try:
             agent_trajectory = self._convert_bedrock_traces_to_messages(traces)
             logger.info(
-                "Bedrock agent produced %d trajectory messages",
-                len(agent_trajectory),
-            )
-
-            return AgentInvocationResponse(
-                agent_trajectory=agent_trajectory,
+                "Bedrock agent produced %d trajectory messages", len(agent_trajectory)
             )
-
-        except (ClientError, ValueError, KeyError, json.JSONDecodeError):
-            logger.exception("Error converting Bedrock traces to messages")
-            return AgentInvocationResponse()
+            return InvokeOutput(agent_trajectory=agent_trajectory)
+        except Exception:
+            logger.exception("Error invoking Bedrock agent")
+            return InvokeOutput()