quraite 0.0.1__py3-none-any.whl

quraite/adapters/langgraph_adapter.py

@@ -0,0 +1,304 @@
+import json
+import uuid
+from typing import Optional
+
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
+from langgraph.graph.state import CompiledStateGraph
+from opentelemetry.trace import TracerProvider
+
+from quraite.adapters.base import BaseAdapter
+from quraite.logger import get_logger
+from quraite.schema.message import AgentMessage, AssistantMessage, MessageContentText
+from quraite.schema.message import SystemMessage as QuraiteSystemMessage
+from quraite.schema.message import ToolCall
+from quraite.schema.message import ToolMessage as QuraiteToolMessage
+from quraite.schema.message import UserMessage
+from quraite.schema.response import AgentInvocationResponse
+from quraite.tracing.constants import QURAITE_ADAPTER_TRACE_PREFIX, Framework
+from quraite.tracing.trace import AgentSpan, AgentTrace
+
+LangchainMessage = HumanMessage | SystemMessage | AIMessage | ToolMessage
+
+logger = get_logger(__name__)
+
+
+class LanggraphAdapter(BaseAdapter):
+    """
+    LangGraph adapter wrapper that converts any LangGraph agent
+    to a standardized callable interface (invoke) and converts the output to List[AgentMessage].
+
+    This class wraps any LangGraph CompiledGraph and provides:
+    - Synchronous invocation via invoke()
+    - Asynchronous invocation via ainvoke()
+    - Automatic conversion to List[AgentMessage] format
+    """
+
+    def __init__(
+        self,
+        agent_graph: CompiledStateGraph,
+        agent_name: str = "LangGraph Agent",
+        tracer_provider: Optional[TracerProvider] = None,
+    ):
+        """
+        Initialize with a pre-configured LangGraph agent
+
+        Args:
+            agent_graph: Any CompiledGraph from LangGraph (must have invoke/ainvoke methods)
+            agent_name: Name of the agent for trajectory metadata
+        """
+        logger.debug("Initializing LanggraphAdapter with agent_name=%s", agent_name)
+        self.agent_graph = agent_graph
+        self.agent_name = agent_name
+        self._init_tracing(tracer_provider, required=False)
+        logger.info(
+            "LanggraphAdapter initialized successfully (tracing_enabled=%s)",
+            bool(tracer_provider),
+        )
+
+    def _prepare_input(
+        self, input: list[AgentMessage]
+    ) -> dict[str, list[HumanMessage]]:
+        """
+        Prepare input for LangGraph agent from List[Message].
+
+        Args:
+            input: List[AgentMessage] containing user_message
+
+        Returns:
+            Dictionary with 'messages' key containing the prepared UserMessage
+        """
+        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 (input_length=%d)", len(input)
+            )
+            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 = next(
+            (
+                content_item.text
+                for content_item in last_user_message.content
+                if content_item.type == "text" and content_item.text
+            ),
+            None,
+        )
+
+        if not text_content:
+            logger.error("No text content found in user message")
+            raise ValueError("No text content found in user message")
+
+        logger.debug("Prepared input with text_content length=%d", len(text_content))
+        return {"messages": [HumanMessage(content=text_content)]}
+
+    def _convert_langchain_messages_to_quraite_messages(
+        self,
+        messages: list[LangchainMessage],
+    ) -> list[AgentMessage]:
+        logger.debug(
+            "Converting %d langchain messages to quraite format", len(messages)
+        )
+        converted_messages: list[AgentMessage] = []
+
+        for idx, msg in enumerate(messages):
+            match msg:
+                case SystemMessage():
+                    logger.debug("Converting SystemMessage at index %d", idx)
+                    converted_messages.append(
+                        QuraiteSystemMessage(
+                            content=[MessageContentText(type="text", text=msg.content)]
+                        )
+                    )
+
+                case HumanMessage():
+                    logger.debug("Converting HumanMessage at index %d", idx)
+                    converted_messages.append(
+                        UserMessage(
+                            content=[MessageContentText(type="text", text=msg.content)]
+                        )
+                    )
+
+                case AIMessage():
+                    logger.debug(
+                        "Converting AIMessage at index %d (has_tool_calls=%s)",
+                        idx,
+                        bool(msg.tool_calls),
+                    )
+                    text_content, tool_calls = self._extract_ai_message_content(msg)
+                    converted_messages.append(
+                        AssistantMessage(
+                            content=text_content if text_content else None,
+                            tool_calls=tool_calls if tool_calls else None,
+                        )
+                    )
+
+                case ToolMessage():
+                    logger.debug(
+                        "Converting ToolMessage at index %d (tool_call_id=%s)",
+                        idx,
+                        msg.tool_call_id,
+                    )
+                    if not msg.content:
+                        tool_message_content = ""
+                    elif isinstance(msg.content, str):
+                        tool_message_content = msg.content
+                    else:
+                        tool_message_content = json.dumps(msg.content)
+
+                    converted_messages.append(
+                        QuraiteToolMessage(
+                            tool_call_id=msg.tool_call_id,
+                            content=[
+                                MessageContentText(
+                                    type="text", text=tool_message_content
+                                )
+                            ],
+                        )
+                    )
+
+        logger.info("Converted %d messages successfully", len(converted_messages))
+        return converted_messages
+
+    def _extract_ai_message_content(
+        self, msg: AIMessage
+    ) -> tuple[list[MessageContentText], list[ToolCall]]:
+        text_content = []
+
+        if msg.content:
+            match msg.content:
+                case str(text):
+                    text_content.append(MessageContentText(type="text", text=text))
+                case list():
+                    text_content.extend(
+                        MessageContentText(type="text", text=content.get("text"))
+                        for content in msg.content
+                        if isinstance(content, dict) and content.get("type") == "text"
+                    )
+
+        tool_calls = []
+        if msg.tool_calls:
+            logger.debug("Extracting %d tool calls from AIMessage", len(msg.tool_calls))
+            tool_calls.extend(
+                ToolCall(
+                    id=tool_call.get("id"),  # type: ignore[union-attr]
+                    name=tool_call.get("name"),  # type: ignore[union-attr]
+                    arguments=tool_call.get("args"),  # type: ignore[union-attr]
+                )
+                for tool_call in msg.tool_calls
+            )
+
+        return text_content, tool_calls
+
+    async def ainvoke(
+        self,
+        input: list[AgentMessage],
+        session_id: str | None,
+    ) -> AgentInvocationResponse:
+        """
+        Asynchronous invocation method - invokes the LangGraph agent and converts the output to List[AgentMessage]
+
+        Args:
+            input: List[AgentMessage] containing user_message
+            session_id: Optional conversation ID for maintaining context
+
+        Returns:
+            List[AgentMessage] or AgentTrace - converted messages from the agent's response or trace
+        """
+        logger.info(
+            "ainvoke called (session_id=%s, input_messages=%d)", session_id, len(input)
+        )
+
+        try:
+            agent_input = self._prepare_input(input)
+            config = {"configurable": {"thread_id": session_id}} if session_id else {}
+
+            if self.tracer_provider:
+                logger.debug("Invoking with tracing enabled")
+                return await self._ainvoke_with_tracing(agent_input, config)
+
+            logger.debug("Invoking without tracing")
+            return await self._ainvoke_without_tracing(agent_input, config)
+
+        except ValueError:
+            logger.exception("ValueError during ainvoke")
+            raise
+        except Exception:
+            logger.exception("Unexpected error during ainvoke")
+            raise
+
+    async def _ainvoke_with_tracing(
+        self,
+        agent_input: dict[str, list[HumanMessage]],
+        config: dict,
+    ) -> AgentInvocationResponse:
+        """Execute ainvoke with tracing enabled."""
+        adapter_trace_id = f"{QURAITE_ADAPTER_TRACE_PREFIX}-{uuid.uuid4()}"
+        logger.debug("Starting traced invocation (trace_id=%s)", adapter_trace_id)
+
+        with self.tracer.start_as_current_span(name=adapter_trace_id):
+            result = await self.agent_graph.ainvoke(agent_input, config=config)
+
+        trace_readable_spans = self.quraite_span_exporter.get_trace_by_testcase(
+            adapter_trace_id
+        )
+
+        if trace_readable_spans:
+            logger.info("Retrieved %d spans from trace", len(trace_readable_spans))
+            agent_trace = AgentTrace(
+                spans=[
+                    AgentSpan.from_readable_oi_span(span)
+                    for span in trace_readable_spans
+                ]
+            )
+
+            trajectory = agent_trace.to_agent_trajectory(framework=Framework.LANGGRAPH)
+            logger.debug("Generated trajectory with %d messages", len(trajectory))
+
+            return AgentInvocationResponse(
+                agent_trace=agent_trace,
+                agent_trajectory=trajectory,
+            )
+
+        logger.warning("No trace spans found for trace_id=%s", adapter_trace_id)
+        return AgentInvocationResponse()
+
+    async def _ainvoke_without_tracing(
+        self,
+        agent_input: dict[str, list[HumanMessage]],
+        config: dict,
+    ) -> AgentInvocationResponse:
+        """Execute ainvoke without tracing."""
+        logger.debug("Starting non-traced invocation")
+        agent_messages = []
+
+        try:
+            async for event in self.agent_graph.astream(agent_input, config=config):
+                logger.debug(
+                    "Received stream event with %d values", len(event.values())
+                )
+                for result in event.values():
+                    if messages := result.get("messages"):
+                        logger.debug("Processing %d messages from event", len(messages))
+                        agent_messages.extend(messages)
+
+            logger.info(
+                "Streaming complete, received %d total messages", len(agent_messages)
+            )
+
+            agent_trajectory = self._convert_langchain_messages_to_quraite_messages(
+                agent_messages
+            )
+
+            return AgentInvocationResponse(
+                agent_trajectory=agent_trajectory,
+            )
+
+        except ValueError:
+            logger.exception("Error converting messages to List[AgentMessage]")
+            return AgentInvocationResponse()

quraite/adapters/langgraph_server_adapter.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+from typing import Annotated, Any, List, Optional, Union
+
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
+from langgraph.pregel.remote import RemoteGraph
+from langgraph_sdk import get_client, get_sync_client
+from pydantic import Discriminator
+
+from quraite.adapters.base import BaseAdapter
+from quraite.logger import get_logger
+from quraite.schema.message import AgentMessage, AssistantMessage, MessageContentText
+from quraite.schema.message import SystemMessage as QuraiteSystemMessage
+from quraite.schema.message import ToolCall, ToolMessage, UserMessage
+from quraite.schema.response import AgentInvocationResponse
+
+LangchainMessage = Annotated[
+    Union[HumanMessage, SystemMessage, AIMessage, ToolMessage],
+    Discriminator(discriminator="type"),
+]
+
+logger = get_logger(__name__)
+
+
+class LanggraphServerAdapter(BaseAdapter):
+    """Remote LangGraph server adapter based on langgraph-sdk.
+
+    Args:
+        base_url: The base URL of the LangGraph server
+        assistant_id: The ID of the assistant to invoke
+        **kwargs: Additional keyword arguments passed directly to
+                 langgraph_sdk.get_client() and get_sync_client().
+                 Common options include:
+                 - api_key: API key for authentication
+                 - headers: Additional HTTP headers
+                 - timeout: Request timeout configuration
+    """
+
+    def __init__(
+        self,
+        *,
+        base_url: str,
+        assistant_id: Optional[str] = None,
+        graph_name: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        self.base_url = base_url
+        self.assistant_id = assistant_id
+        self.graph_name = graph_name
+
+        logger.debug(
+            "Initializing LanggraphServerAdapter (base_url=%s, assistant_id=%s, graph_name=%s)",
+            base_url,
+            assistant_id,
+            graph_name,
+        )
+        try:
+            sync_client = get_sync_client(url=self.base_url, **kwargs)
+            async_client = get_client(url=self.base_url, **kwargs)
+            if self.assistant_id:
+                self.remote_graph = RemoteGraph(
+                    self.assistant_id,
+                    url=self.base_url,
+                    sync_client=sync_client,
+                    client=async_client,
+                )
+            else:
+                self.remote_graph = RemoteGraph(
+                    self.graph_name,
+                    url=self.base_url,
+                    sync_client=sync_client,
+                    client=async_client,
+                )
+        except Exception as exc:
+            raise RuntimeError(
+                f"Failed to initialize LangGraph RemoteGraph for {self.base_url}: {exc}"
+            )
+        logger.info(
+            "LanggraphServerAdapter initialized (assistant_id=%s, graph_name=%s)",
+            self.assistant_id,
+            self.graph_name,
+        )
+
+    def _prepare_input(self, input: List[AgentMessage]) -> Any:
+        """
+        Prepare input for LangGraph agent from List[AgentMessage].
+
+        Args:
+            input: List[AgentMessage] containing user_message
+
+        Returns:
+            Dict with messages list containing user_message
+        """
+        logger.debug("Preparing Langgraph server input from %d messages", len(input))
+        if not input or input[-1].role != "user":
+            logger.error("Langgraph server input missing user message")
+            raise ValueError("No user message found in the input")
+
+        last_user_message = input[-1]
+        # Check if content list is not empty and has text
+        if not last_user_message.content:
+            logger.error("Langgraph server user message missing content")
+            raise ValueError("User message has no content")
+
+        # Find the first text content item
+        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("Langgraph server user message missing text content")
+            raise ValueError("No text content found in user message")
+
+        logger.debug(
+            "Prepared Langgraph server input (text_length=%d)", len(text_content)
+        )
+        return {"messages": [HumanMessage(content=text_content).model_dump()]}
+
+    def _convert_langchain_messages_to_quraite_messages(
+        self,
+        messages: List[dict],
+    ) -> List[AgentMessage]:
+        logger.debug(
+            "Converting %d Langgraph server messages to quraite format", len(messages)
+        )
+        converted_messages: List[AgentMessage] = []
+
+        for msg in messages:
+            if msg.get("type") == "system":
+                converted_messages.append(
+                    QuraiteSystemMessage(
+                        content=[
+                            MessageContentText(type="text", text=msg.get("content", ""))
+                        ],
+                    )
+                )
+
+            elif msg.get("type") == "human":
+                converted_messages.append(
+                    UserMessage(
+                        content=[
+                            MessageContentText(type="text", text=msg.get("content", ""))
+                        ],
+                    )
+                )
+
+            elif msg.get("type") == "ai":
+                text_content: List[MessageContentText] = []
+                tool_calls_list: List[ToolCall] = []
+
+                # Extract text content - sometimes it's a string, sometimes a list of dicts
+                content = msg.get("content")
+                if isinstance(content, str) and content:
+                    text_content.append(MessageContentText(type="text", text=content))
+                elif isinstance(content, list):
+                    for content_item in content:
+                        if isinstance(content_item, dict):
+                            if content_item.get("type") == "text" and content_item.get(
+                                "text"
+                            ):
+                                text_content.append(
+                                    MessageContentText(
+                                        type="text", text=content_item.get("text")
+                                    )
+                                )
+
+                # Extract tool calls if present
+                if msg.get("tool_calls"):
+                    for tool_call in msg.get("tool_calls"):
+                        if isinstance(tool_call, dict):
+                            tool_calls_list.append(
+                                ToolCall(
+                                    id=tool_call.get("id", ""),
+                                    name=tool_call.get("name", ""),
+                                    arguments=tool_call.get("args", {}),
+                                )
+                            )
+
+                converted_messages.append(
+                    AssistantMessage(
+                        content=text_content if text_content else None,
+                        tool_calls=tool_calls_list if tool_calls_list else None,
+                    )
+                )
+
+            elif msg.get("type") == "tool":
+                tool_content = msg.get("content", "")
+                converted_messages.append(
+                    ToolMessage(
+                        tool_call_id=msg.get("tool_call_id", ""),
+                        content=[
+                            MessageContentText(type="text", text=str(tool_content))
+                        ],
+                    )
+                )
+
+            else:
+                # Skip unsupported message types
+                continue
+
+        logger.info(
+            "Langgraph server message conversion produced %d messages",
+            len(converted_messages),
+        )
+        return converted_messages
+
+    async def ainvoke(
+        self,
+        input: List[AgentMessage],
+        session_id: Annotated[Union[str, None], "Thread ID used by LangGraph API"],
+    ) -> AgentInvocationResponse:
+        agent_messages = []
+        agent_input = self._prepare_input(input)
+        if session_id:
+            config = {"configurable": {"thread_id": session_id}}
+        else:
+            config = {}
+
+        try:
+            logger.info("Langgraph server ainvoke called (session_id=%s)", session_id)
+            async for event in self.remote_graph.astream(agent_input, config=config):
+                for _, result in event.items():
+                    if result.get("messages"):
+                        logger.debug(
+                            "Langgraph server received %d messages from stream chunk",
+                            len(result.get("messages")),
+                        )
+                        agent_messages += result.get("messages")
+
+        except Exception as e:
+            logger.exception("Error invoking Langgraph remote graph")
+            raise RuntimeError(f"Error invoking LangGraph agent: {e}") from e
+
+        try:
+            # Convert to List[AgentMessage]
+            agent_trajectory = self._convert_langchain_messages_to_quraite_messages(
+                agent_messages
+            )
+            logger.info(
+                "Langgraph server ainvoke produced %d trajectory messages",
+                len(agent_trajectory),
+            )
+
+            return AgentInvocationResponse(
+                agent_trajectory=agent_trajectory,
+            )
+
+        except ValueError:
+            logger.exception("Langgraph server conversion to AgentMessage failed")
+            return AgentInvocationResponse()