quraite 0.0.1__py3-none-any.whl

quraite/adapters/pydantic_ai_adapter.py

@@ -0,0 +1,312 @@
+from typing import Any, List, Union
+
+from opentelemetry.trace import TracerProvider
+from pydantic_ai import Agent
+from pydantic_ai.messages import (
+    BuiltinToolCallPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+
+from quraite.adapters.base import BaseAdapter
+from quraite.logger import get_logger
+from quraite.schema.message import (
+    AgentMessage,
+    AssistantMessage,
+    MessageContentText,
+    ToolCall,
+    ToolMessage,
+)
+from quraite.schema.response import AgentInvocationResponse
+from quraite.tracing.constants import QURAITE_ADAPTER_TRACE_PREFIX, Framework
+from quraite.tracing.trace import AgentSpan, AgentTrace
+
+logger = get_logger(__name__)
+
+
+class PydanticAIAdapter(BaseAdapter):
+    """
+    Pydantic AI adapter wrapper that converts any Pydantic AI agent
+    to a standardized callable interface (invoke) and converts the output to List[AgentMessage].
+
+    This class wraps any Pydantic AI Agent and provides:
+    - Synchronous invocation via invoke()
+    - Asynchronous invocation via ainvoke()
+    - Automatic conversion to List[AgentMessage] format
+    - Access to message history and traces
+    """
+
+    def __init__(
+        self,
+        agent: Agent,
+        agent_name: str = "Pydantic AI Agent",
+        tracer_provider: TracerProvider | None = None,
+    ):
+        """
+        Initialize with a pre-configured Pydantic AI agent
+
+        Args:
+            agent: A Pydantic AI Agent instance
+            agent_name: Name of the agent for trajectory metadata
+        """
+        self.agent = agent
+        self.agent_name = agent_name
+        # Store session state for conversation context
+        self._sessions: dict[str, Any] = {}
+        self._init_tracing(tracer_provider, required=False)
+        logger.info(
+            "PydanticAIAdapter initialized (agent_name=%s, tracing_enabled=%s)",
+            agent_name,
+            bool(tracer_provider),
+        )
+
+    def _convert_pydantic_ai_messages_to_messages(
+        self, messages: List[Any]
+    ) -> List[AgentMessage]:
+        """
+        Convert Pydantic AI ModelMessage objects (with parts) to SDK Message format.
+
+        Args:
+            messages: List of Pydantic AI ModelMessage objects
+
+        Returns:
+            List[AgentMessage]: Converted messages in SDK format
+        """
+        converted_messages: List[AgentMessage] = []
+
+        for msg in messages:
+            if not hasattr(msg, "parts"):
+                continue
+
+            content: List[MessageContentText] = []
+            tool_calls_list: List[ToolCall] = []
+
+            for part in msg.parts:
+                if isinstance(part, SystemPromptPart) or isinstance(
+                    part, UserPromptPart
+                ):
+                    continue
+
+                elif isinstance(part, TextPart):
+                    if part.content:
+                        content.append(
+                            MessageContentText(type="text", text=part.content)
+                        )
+                elif isinstance(part, ToolCallPart) or isinstance(
+                    part, BuiltinToolCallPart
+                ):
+                    tool_call_id = part.tool_call_id
+                    tool_name = part.tool_name
+                    args = part.args
+
+                    tool_calls_list.append(
+                        ToolCall(
+                            id=tool_call_id,
+                            name=tool_name,
+                            arguments=args if isinstance(args, dict) else {},
+                        )
+                    )
+
+                elif isinstance(part, ToolReturnPart):
+                    if content or tool_calls_list:
+                        converted_messages.append(
+                            AssistantMessage(
+                                content=content if content else None,
+                                tool_calls=tool_calls_list if tool_calls_list else None,
+                            )
+                        )
+                        content = []
+                        tool_calls_list = []
+
+                    tool_call_id = part.tool_call_id
+
+                    converted_messages.append(
+                        ToolMessage(
+                            tool_name=part.tool_name,
+                            tool_call_id=tool_call_id,
+                            content=[
+                                MessageContentText(type="text", text=str(part.content))
+                            ],
+                        )
+                    )
+
+            if content or tool_calls_list:
+                converted_messages.append(
+                    AssistantMessage(
+                        content=content if content else None,
+                        tool_calls=tool_calls_list if tool_calls_list else None,
+                    )
+                )
+
+        logger.info(
+            "Converted %d Pydantic AI messages to agent messages",
+            len(converted_messages),
+        )
+        return converted_messages
+
+    def _prepare_input(self, input: List[AgentMessage]) -> str:
+        """
+        Prepare input for Pydantic AI agent from List[AgentMessage].
+
+        Args:
+            input: List[AgentMessage] containing user_message
+
+        Returns:
+            str: User message text
+        """
+        logger.debug("Preparing Pydantic AI input from %d messages", len(input))
+        if not input or input[-1].role != "user":
+            logger.error("Pydantic AI input missing user message")
+            raise ValueError("No user message found in the input")
+
+        last_user_message = input[-1]
+        if not last_user_message.content:
+            logger.error("Pydantic AI user message missing 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("Pydantic AI user message missing text content")
+            raise ValueError("No text content found in user message")
+
+        logger.debug("Prepared Pydantic AI input (text_length=%d)", len(text_content))
+        return text_content
+
+    async def ainvoke(
+        self,
+        input: List[AgentMessage],
+        session_id: Union[str, None],
+    ) -> AgentInvocationResponse:
+        """
+        Asynchronous invocation method - invokes the Pydantic AI agent and converts the output to List[AgentMessage]
+
+        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(
+            "Pydantic AI ainvoke called (session_id=%s, input_messages=%d)",
+            session_id,
+            len(input),
+        )
+        agent_input = self._prepare_input(input)
+        session_id = session_id or "default"
+
+        try:
+            # Get message history for this session (for multi-turn conversations)
+            message_history = self._sessions.get(session_id, None)
+
+            # Run with or without tracing depending on configuration
+            if self.tracer_provider:
+                return await self._ainvoke_with_tracing(
+                    agent_input, session_id, message_history
+                )
+
+            return await self._ainvoke_without_tracing(
+                agent_input, session_id, message_history
+            )
+
+        except Exception as e:
+            logger.exception("Error invoking Pydantic AI agent")
+            raise RuntimeError(f"Error invoking Pydantic AI agent: {e}") from e
+
+    async def _ainvoke_with_tracing(
+        self,
+        agent_input: str,
+        session_id: str,
+        message_history: Any | None,
+    ) -> AgentInvocationResponse:
+        """Execute ainvoke with tracing enabled."""
+        from uuid import uuid4
+
+        adapter_trace_id = f"{QURAITE_ADAPTER_TRACE_PREFIX}-{uuid4()}"
+        logger.debug(
+            "Starting Pydantic AI traced invocation (trace_id=%s, session_id=%s)",
+            adapter_trace_id,
+            session_id,
+        )
+
+        with self.tracer.start_as_current_span(name=adapter_trace_id):
+            # Run the agent asynchronously with message history for context
+            if message_history:
+                result = await self.agent.run(
+                    agent_input,
+                    message_history=message_history,
+                )
+            else:
+                result = await self.agent.run(agent_input)
+
+        # Store the complete updated message history for this session
+        self._sessions[session_id] = result.all_messages()
+
+        trace_readable_spans = self.quraite_span_exporter.get_trace_by_testcase(
+            adapter_trace_id
+        )
+
+        if trace_readable_spans:
+            agent_trace = AgentTrace(
+                spans=[
+                    AgentSpan.from_readable_oi_span(span)
+                    for span in trace_readable_spans
+                ]
+            )
+            logger.info(
+                "Pydantic AI trace collected %d spans for trace_id=%s",
+                len(trace_readable_spans),
+                adapter_trace_id,
+            )
+        else:
+            logger.warning(
+                "No spans exported for Pydantic AI trace_id=%s", adapter_trace_id
+            )
+
+        return AgentInvocationResponse(
+            agent_trace=agent_trace,
+            agent_trajectory=agent_trace.to_agent_trajectory(
+                framework=Framework.PYDANTIC_AI
+            ),
+        )
+
+    async def _ainvoke_without_tracing(
+        self,
+        agent_input: str,
+        session_id: str,
+        message_history: Any | None,
+    ) -> AgentInvocationResponse:
+        """Execute ainvoke without tracing."""
+        # Run the agent asynchronously with message history for context
+        if message_history:
+            result = await self.agent.run(
+                agent_input,
+                message_history=message_history,
+            )
+        else:
+            result = await self.agent.run(agent_input)
+
+        # Store the complete updated message history for this session
+        self._sessions[session_id] = result.all_messages()
+
+        # Convert only the NEW messages to SDK Message format
+        agent_trajectory = self._convert_pydantic_ai_messages_to_messages(
+            result.new_messages()
+        )
+
+        logger.info(
+            "Pydantic AI produced %d trajectory messages without tracing",
+            len(agent_trajectory),
+        )
+
+        return AgentInvocationResponse(
+            agent_trajectory=agent_trajectory,
+        )

quraite/adapters/smolagents_adapter.py

@@ -0,0 +1,152 @@
+import asyncio
+import uuid
+from typing import List, Optional
+
+from opentelemetry.trace import TracerProvider
+from smolagents import CodeAgent
+
+from quraite.adapters.base import BaseAdapter
+from quraite.logger import get_logger
+from quraite.schema.message import AgentMessage
+from quraite.schema.response import AgentInvocationResponse
+from quraite.tracing.constants import QURAITE_ADAPTER_TRACE_PREFIX, Framework
+from quraite.tracing.trace import AgentSpan, AgentTrace
+
+logger = get_logger(__name__)
+
+
+class SmolagentsAdapter(BaseAdapter):
+    """
+    Smolagents adapter wrapper that converts any Smolagents CodeAgent
+    to a standardized callable interface (ainvoke) with tracing support.
+
+    This class wraps any CodeAgent and provides:
+    - Asynchronous invocation via ainvoke()
+    - OpenTelemetry tracing integration
+    - Required tracing (returns AgentTrace containing spans)
+    """
+
+    def __init__(
+        self,
+        agent: CodeAgent,
+        agent_name: str = "Smolagents Agent",
+        tracer_provider: Optional[TracerProvider] = None,
+    ):
+        """
+        Initialize with a pre-configured Smolagents CodeAgent
+
+        Args:
+            agent: A Smolagents CodeAgent instance
+            agent_name: Name of the agent for trajectory metadata
+            tracer_provider: TracerProvider for tracing (required)
+        """
+        logger.debug(
+            "Initializing SmolagentsAdapter with agent_name=%s (tracing_required=True)",
+            agent_name,
+        )
+        self._init_tracing(tracer_provider, required=True)
+
+        self.agent = agent
+        self.agent_name = agent_name
+        logger.info("SmolagentsAdapter initialized successfully")
+
+    def _prepare_input(self, input: List[AgentMessage]) -> str:
+        """
+        Prepare input for Smolagents CodeAgent from List[AgentMessage].
+
+        Args:
+            input: List[AgentMessage] containing user_message
+
+        Returns:
+            str: User message text
+        """
+        logger.debug("Preparing Smolagents input from %d messages", len(input))
+        if not input or input[-1].role != "user":
+            logger.error("Smolagents input missing user message")
+            raise ValueError("No user message found in the input")
+
+        last_user_message = input[-1]
+        if not last_user_message.content:
+            logger.error("Smolagents user message missing 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("Smolagents user message missing text content")
+            raise ValueError("No text content found in user message")
+
+        logger.debug("Prepared Smolagents input (text_length=%d)", len(text_content))
+        return text_content
+
+    async def ainvoke(
+        self,
+        input: List[AgentMessage],
+        session_id: Optional[str] = None,
+    ) -> AgentInvocationResponse:
+        """
+        Asynchronous invocation method - invokes the Smolagents CodeAgent with tracing
+
+        Args:
+            input: List[AgentMessage] containing user_message
+            session_id: Optional conversation ID (for parity with other adapters)
+
+        Returns:
+            AgentTrace - trace with spans captured during invocation
+        """
+        _ = session_id  # Currently unused but kept for interface compatibility
+        logger.info(
+            "Smolagents ainvoke called (session_id=%s, input_messages=%d)",
+            session_id,
+            len(input),
+        )
+        agent_input = self._prepare_input(input)
+
+        try:
+            return await self._ainvoke_with_tracing(agent_input)
+        except Exception as exc:
+            logger.exception("Error invoking Smolagents agent")
+            raise RuntimeError(f"Error invoking Smolagents agent: {exc}") from exc
+
+    async def _ainvoke_with_tracing(self, agent_input: str) -> AgentInvocationResponse:
+        """Execute ainvoke with tracing enabled."""
+        adapter_trace_id = f"{QURAITE_ADAPTER_TRACE_PREFIX}-{uuid.uuid4()}"
+        logger.debug(
+            "Starting Smolagents traced invocation (trace_id=%s)", adapter_trace_id
+        )
+
+        with self.tracer.start_as_current_span(name=adapter_trace_id):
+            # Run the agent synchronously inside a thread to avoid blocking
+            await asyncio.to_thread(self.agent.run, agent_input)
+
+        trace_readable_spans = self.quraite_span_exporter.get_trace_by_testcase(
+            adapter_trace_id
+        )
+
+        if trace_readable_spans:
+            logger.info(
+                "Smolagents trace collected %d spans for trace_id=%s",
+                len(trace_readable_spans),
+                adapter_trace_id,
+            )
+            agent_trace = AgentTrace(
+                spans=[
+                    AgentSpan.from_readable_oi_span(span)
+                    for span in trace_readable_spans
+                ],
+            )
+        else:
+            logger.warning(
+                "No spans captured for Smolagents trace_id=%s", adapter_trace_id
+            )
+
+        return AgentInvocationResponse(
+            agent_trace=agent_trace,
+            agent_trajectory=agent_trace.to_agent_trajectory(
+                framework=Framework.SMOLAGENTS
+            ),
+        )

quraite/logger.py

@@ -0,0 +1,62 @@
+import logging
+import os
+
+_logger = logging.getLogger("quraite")
+_logger.addHandler(logging.NullHandler())
+
+
+def set_log_level(level: int | str) -> None:
+    """
+    Configure the root quraite logger programmatically.
+
+    Args:
+        level: Either a logging level integer or its string name (e.g. "INFO").
+    """
+    resolved_level: int
+    if isinstance(level, str):
+        try:
+            resolved_level = getattr(logging, level.upper())
+        except AttributeError as exc:
+            raise ValueError(f"Invalid log level: {level}") from exc
+    else:
+        resolved_level = level
+
+    _logger.setLevel(resolved_level)
+
+    # Ensure a stream handler is attached so logs are emitted somewhere useful.
+    if not any(
+        isinstance(handler, logging.StreamHandler) for handler in _logger.handlers
+    ):
+        handler = logging.StreamHandler()
+        handler.setFormatter(
+            logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+        )
+        _logger.addHandler(handler)
+
+
+def get_logger(name: str | None = None) -> logging.Logger:
+    """
+    Get a logger for quraite submodules.
+
+    Args:
+        name: Module name (typically __name__)
+
+    Returns:
+        A logging.Logger instance for the quraite package.
+    """
+    if name is None:
+        return _logger
+
+    if name.startswith("quraite"):
+        return logging.getLogger(name)
+
+    return logging.getLogger(f"quraite.{name}")
+
+
+# Set level from environment variable if provided
+_env_level = os.getenv("LOG_LEVEL")
+if _env_level:
+    try:
+        set_log_level(_env_level)
+    except ValueError:
+        pass

quraite/schema/message.py

@@ -0,0 +1,54 @@
+from typing import Any, List, Literal, Optional, TypeAlias, Union
+
+from pydantic import BaseModel
+
+
+class MessageContentText(BaseModel):
+    type: Literal["text"] = "text"
+    text: str
+
+
+class MessageContentReasoning(BaseModel):
+    type: Literal["reasoning"] = "reasoning"
+    reasoning: str
+
+
+class UserMessage(BaseModel):
+    role: Literal["user"] = "user"
+    name: Optional[str] = None
+    content: List[MessageContentText]
+
+
+class DeveloperMessage(BaseModel):
+    role: Literal["developer"] = "developer"
+    content: List[MessageContentText]
+
+
+class SystemMessage(BaseModel):
+    role: Literal["system"] = "system"
+    content: List[MessageContentText]
+
+
+class ToolCall(BaseModel):
+    id: str
+    name: str
+    arguments: dict[str, Any]
+
+
+class AssistantMessage(BaseModel):
+    role: Literal["assistant"] = "assistant"
+    agent_name: Optional[str] = None
+    content: Optional[List[Union[MessageContentText, MessageContentReasoning]]] = None
+    tool_calls: Optional[List[ToolCall]] = None
+
+
+class ToolMessage(BaseModel):
+    role: Literal["tool"] = "tool"
+    tool_name: Optional[str] = None
+    tool_call_id: Optional[str] = None
+    content: List[MessageContentText]
+
+
+AgentMessage: TypeAlias = Union[
+    UserMessage, DeveloperMessage, SystemMessage, AssistantMessage, ToolMessage
+]

quraite/schema/response.py

@@ -0,0 +1,16 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+from quraite.schema.message import AgentMessage
+from quraite.tracing.trace import AgentTrace
+
+
+class AgentInvocationResponse(BaseModel):
+    """
+    Response model for agent invocation.
+    """
+
+    agent_trace: Optional[AgentTrace] = None
+    agent_trajectory: Optional[List[AgentMessage]] = None
+    agent_final_response: Optional[str] = None

quraite/serve/__init__.py

@@ -0,0 +1 @@
+