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

quraite/adapters/langchain_adapter.py

@@ -1,19 +1,24 @@
 import json
-from typing import Optional
 
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.runnables import RunnableConfig
 from langgraph.graph.state import CompiledStateGraph
-from opentelemetry.trace import TracerProvider
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider
 
 from quraite.adapters.base import BaseAdapter
+from quraite.constants.framework import Framework
 from quraite.logger import get_logger
-from quraite.schema.message import AgentMessage, AssistantMessage, MessageContentText
+from quraite.schema.invoke import InvokeInput, InvokeOutput
+from quraite.schema.message import (
+    AgentMessage,
+    AssistantMessage,
+    MessageContentText,
+    ToolCall,
+    UserMessage,
+)
 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 Framework
 from quraite.tracing.trace import AgentSpan, AgentTrace
 
 LangchainMessage = HumanMessage | SystemMessage | AIMessage | ToolMessage
@@ -22,84 +27,33 @@ logger = get_logger(__name__)
 
 
 class LangchainAdapter(BaseAdapter):
-    """
-    LangChain adapter wrapper that converts any LangChain agent
-    to a standardized callable interface (invoke) and converts the output to List[AgentMessage].
-
-    This class wraps any LangChain CompiledGraph and provides:
-    - Synchronous invocation via invoke()
-    - Asynchronous invocation via ainvoke()
-    - Automatic conversion to List[AgentMessage] format
-    """
+    """LangChain adapter for CompiledStateGraph agents."""
 
     def __init__(
         self,
         agent_graph: CompiledStateGraph,
+        *,
+        tracer_provider: SDKTracerProvider | None = None,
         agent_name: str = "LangChain Agent",
-        tracer_provider: Optional[TracerProvider] = None,
     ):
         """
-        Initialize with a pre-configured LangChain agent
+        Initialize LangChain adapter.
 
         Args:
-            agent_graph: Any CompiledGraph from LangChain (must have invoke/ainvoke methods)
-            agent_name: Name of the agent for trajectory metadata
+            agent_graph: LangChain CompiledStateGraph
+            tracer_provider: Optional TracerProvider from setup_tracing()
+            agent_name: Agent name for metadata
         """
-        logger.debug("Initializing LangchainAdapter 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(
-            "LangchainAdapter initialized successfully (tracing_enabled=%s)",
-            bool(tracer_provider),
-        )
-
-    def _prepare_input(
-        self, input: list[AgentMessage]
-    ) -> dict[str, list[HumanMessage]]:
-        """
-        Prepare input for LangChain 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)]}
+        self._init_tracer(tracer_provider)
+        logger.info("LangchainAdapter initialized (tracing=%s)", bool(tracer_provider))
 
     def _convert_langchain_messages_to_quraite_messages(
         self,
         messages: list[LangchainMessage],
     ) -> list[AgentMessage]:
+        """Convert LangChain messages to Quraite format."""
         logger.debug(
             "Converting %d langchain messages to quraite format", len(messages)
         )
@@ -111,7 +65,7 @@ class LangchainAdapter(BaseAdapter):
                     logger.debug("Converting SystemMessage at index %d", idx)
                     converted_messages.append(
                         QuraiteSystemMessage(
-                            content=[MessageContentText(type="text", text=msg.content)]
+                            content=[MessageContentText(type="text", text=msg.content or "")]
                         )
                     )
 
@@ -119,7 +73,7 @@ class LangchainAdapter(BaseAdapter):
                     logger.debug("Converting HumanMessage at index %d", idx)
                     converted_messages.append(
                         UserMessage(
-                            content=[MessageContentText(type="text", text=msg.content)]
+                            content=[MessageContentText(type="text", text=msg.content or "")]
                         )
                     )
 
@@ -132,8 +86,8 @@ class LangchainAdapter(BaseAdapter):
                     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,
+                            content=text_content,
+                            tool_calls=tool_calls,
                         )
                     )
 
@@ -167,6 +121,7 @@ class LangchainAdapter(BaseAdapter):
     def _extract_ai_message_content(
         self, msg: AIMessage
     ) -> tuple[list[MessageContentText], list[ToolCall]]:
+        """Extract text and tool calls from AIMessage."""
         text_content = []
 
         if msg.content:
@@ -175,7 +130,7 @@ class LangchainAdapter(BaseAdapter):
                     text_content.append(MessageContentText(type="text", text=text))
                 case list():
                     text_content.extend(
-                        MessageContentText(type="text", text=content.get("text"))
+                        MessageContentText(type="text", text=content.get("text", ""))
                         for content in msg.content
                         if isinstance(content, dict) and content.get("type") == "text"
                     )
@@ -185,57 +140,35 @@ class LangchainAdapter(BaseAdapter):
             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]
+                    id=tool_call.get("id", ""),     
+                    name=tool_call.get("name", ""),
+                    arguments=tool_call.get("args", {}),  
                 )
                 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 LangChain agent and converts the output to List[AgentMessage]
-
-        Args:
-            input: List[AgentMessage] containing user_message
-            session_id: Optional conversation ID for maintaining context
+    async def ainvoke(self, input: InvokeInput) -> InvokeOutput:
+        """Invoke LangChain agent and return response."""
+        logger.info("ainvoke called (session_id=%s)", input.session_id)
 
-        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)
+        agent_input = {"messages": [HumanMessage(content=input.user_message_str())]}
+        config: RunnableConfig = (
+            {"configurable": {"thread_id": input.session_id}}
+            if input.session_id
+            else {}
         )
 
-        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
+        if self.tracer_provider:
+            return await self._ainvoke_with_tracing(agent_input, config)
+        return await self._ainvoke_without_tracing(agent_input, config)
 
     async def _ainvoke_with_tracing(
         self,
         agent_input: dict[str, list[HumanMessage]],
-        config: dict,
-    ) -> AgentInvocationResponse:
+        config: RunnableConfig,
+    ) -> InvokeOutput:
         """Execute ainvoke with tracing enabled."""
         with self.tracer.start_as_current_span("langchain_invocation") as span:
             trace_id = span.get_span_context().trace_id
@@ -245,7 +178,7 @@ class LangchainAdapter(BaseAdapter):
             )
             _ = await self.agent_graph.ainvoke(agent_input, config=config)
 
-        trace_readable_spans = self.quraite_span_exporter.get_spans_by_trace_id(
+        trace_readable_spans: list[ReadableSpan] = self.quraite_span_exporter.get_spans_by_trace_id(
             trace_id
         )
 
@@ -261,19 +194,19 @@ class LangchainAdapter(BaseAdapter):
             trajectory = agent_trace.to_agent_trajectory(framework=Framework.LANGCHAIN)
             logger.debug("Generated trajectory with %d messages", len(trajectory))
 
-            return AgentInvocationResponse(
+            return InvokeOutput(
                 agent_trace=agent_trace,
                 agent_trajectory=trajectory,
             )
 
         logger.warning("No trace spans found for trace_id=%s", trace_id)
-        return AgentInvocationResponse()
+        return InvokeOutput()
 
     async def _ainvoke_without_tracing(
         self,
         agent_input: dict[str, list[HumanMessage]],
-        config: dict,
-    ) -> AgentInvocationResponse:
+        config: RunnableConfig,
+    ) -> InvokeOutput:
         """Execute ainvoke without tracing."""
         logger.debug("Starting non-traced invocation")
         agent_messages = []
@@ -296,10 +229,10 @@ class LangchainAdapter(BaseAdapter):
                 agent_messages
             )
 
-            return AgentInvocationResponse(
+            return InvokeOutput(
                 agent_trajectory=agent_trajectory,
             )
 
         except ValueError:
             logger.exception("Error converting messages to List[AgentMessage]")
-            return AgentInvocationResponse()
+            return InvokeOutput()

quraite/adapters/langchain_server_adapter.py

@@ -1,23 +1,21 @@
 from __future__ import annotations
 
-from typing import Annotated, Any, List, Optional, Union
-
-from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.messages import HumanMessage, 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.invoke import InvokeInput, InvokeOutput
+from quraite.schema.message import (
+    AgentMessage,
+    AssistantMessage,
+    MessageContentText,
+    ToolCall,
+    ToolMessage,
+    UserMessage,
+)
 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__)
 
@@ -40,8 +38,8 @@ class LangchainServerAdapter(BaseAdapter):
         self,
         *,
         base_url: str,
-        assistant_id: Optional[str] = None,
-        graph_name: Optional[str] = None,
+        assistant_id: str | None = None,
+        graph_name: str | None = None,
         **kwargs,
     ) -> None:
         self.base_url = base_url
@@ -81,51 +79,14 @@ class LangchainServerAdapter(BaseAdapter):
             self.graph_name,
         )
 
-    def _prepare_input(self, input: List[AgentMessage]) -> Any:
-        """
-        Prepare input for LangChain agent from List[AgentMessage].
-
-        Args:
-            input: List[AgentMessage] containing user_message
-
-        Returns:
-            Dict with messages list containing user_message
-        """
-        logger.debug("Preparing Langchain server input from %d messages", len(input))
-        if not input or input[-1].role != "user":
-            logger.error("Langchain 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("Langchain 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("Langchain server user message missing text content")
-            raise ValueError("No text content found in user message")
-
-        logger.debug(
-            "Prepared Langchain 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]:
+        messages: list[dict],
+    ) -> list[AgentMessage]:
         logger.debug(
             "Converting %d Langchain server messages to quraite format", len(messages)
         )
-        converted_messages: List[AgentMessage] = []
+        converted_messages: list[AgentMessage] = []
 
         for msg in messages:
             if msg.get("type") == "system":
@@ -147,8 +108,8 @@ class LangchainServerAdapter(BaseAdapter):
                 )
 
             elif msg.get("type") == "ai":
-                text_content: List[MessageContentText] = []
-                tool_calls_list: List[ToolCall] = []
+                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")
@@ -186,12 +147,12 @@ class LangchainServerAdapter(BaseAdapter):
                 )
 
             elif msg.get("type") == "tool":
-                tool_content = msg.get("content", "")
+                tool_content = msg.get("content", "") or ""
                 converted_messages.append(
                     ToolMessage(
                         tool_call_id=msg.get("tool_call_id", ""),
                         content=[
-                            MessageContentText(type="text", text=str(tool_content))
+                            MessageContentText(type="text", text=str(tool_content or ""))
                         ],
                     )
                 )
@@ -206,47 +167,34 @@ class LangchainServerAdapter(BaseAdapter):
         )
         return converted_messages
 
-    async def ainvoke(
-        self,
-        input: List[AgentMessage],
-        session_id: Annotated[Union[str, None], "Thread ID used by LangChain API"],
-    ) -> AgentInvocationResponse:
-        agent_messages = []
-        agent_input = self._prepare_input(input)
-        if session_id:
-            config = {"configurable": {"thread_id": session_id}}
-        else:
-            config = {}
+    async def ainvoke(self, input: InvokeInput) -> InvokeOutput:
+        """Invoke LangChain server agent and return response."""
+        logger.info("Langchain server ainvoke called (session_id=%s)", input.session_id)
+
+        agent_input = {
+            "messages": [HumanMessage(content=input.user_message_str()).model_dump()]
+        }
+        config = (
+            {"configurable": {"thread_id": input.session_id}}
+            if input.session_id
+            else {}
+        )
 
         try:
-            logger.info("Langchain server ainvoke called (session_id=%s)", session_id)
+            agent_messages = []
             async for event in self.remote_graph.astream(agent_input, config=config):
                 for _, result in event.items():
                     if result.get("messages"):
-                        logger.debug(
-                            "Langchain server received %d messages from stream chunk",
-                            len(result.get("messages")),
-                        )
                         agent_messages += result.get("messages")
 
-        except Exception as e:
-            logger.exception("Error invoking Langchain remote graph")
-            raise RuntimeError(f"Error invoking LangChain agent: {e}") from e
-
-        try:
-            # Convert to List[AgentMessage]
             agent_trajectory = self._convert_langchain_messages_to_quraite_messages(
                 agent_messages
             )
             logger.info(
-                "Langchain server ainvoke produced %d trajectory messages",
+                "Langchain server produced %d trajectory messages",
                 len(agent_trajectory),
             )
-
-            return AgentInvocationResponse(
-                agent_trajectory=agent_trajectory,
-            )
-
-        except ValueError:
-            logger.exception("Langchain server conversion to AgentMessage failed")
-            return AgentInvocationResponse()
+            return InvokeOutput(agent_trajectory=agent_trajectory)
+        except Exception as e:
+            logger.exception("Error invoking Langchain remote graph")
+            raise RuntimeError(f"Error invoking LangChain agent: {e}") from e

quraite/adapters/langflow_adapter.py

@@ -1,12 +1,13 @@
 import asyncio
 import json
 import uuid
-from typing import Any, Dict, List, Union
+from typing import Any
 
 import aiohttp
 
 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,
@@ -14,7 +15,6 @@ from quraite.schema.message import (
     ToolCall,
     ToolMessage,
 )
-from quraite.schema.response import AgentInvocationResponse
 
 logger = get_logger(__name__)
 
@@ -33,8 +33,8 @@ class LangflowAdapter(BaseAdapter):
 
     def _convert_api_output_to_messages(
         self,
-        response: Dict[str, Any],
-    ) -> List[AgentMessage]:
+        response: dict[str, Any],
+    ) -> list[AgentMessage]:
         logger.debug(
             "Converting Langflow response (root_keys=%s)",
             list(response.keys()),
@@ -45,7 +45,7 @@ class LangflowAdapter(BaseAdapter):
         contents = content_blocks[0]["contents"]
 
         # Assume everything sequential.
-        ai_trajectory: List[AgentMessage] = []
+        ai_trajectory: list[AgentMessage] = []
         for step in contents:
             if step["type"] == "text":
                 if step["header"]["title"] == "Input":
@@ -94,35 +94,11 @@ class LangflowAdapter(BaseAdapter):
         )
         return ai_trajectory
 
-    def _prepare_input(self, input: List[AgentMessage]) -> str:
-        logger.debug("Preparing Langflow input from %d messages", len(input))
-        if not input or input[-1].role != "user":
-            logger.error("Langflow 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("Langflow input 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("Langflow input missing text content")
-            raise ValueError("No text content found in user message")
-
-        logger.debug("Prepared Langflow input (text_length=%d)", len(text_content))
-        return text_content
-
     async def _aapi_call(
         self,
         query: str,
         sessionId: str,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         payload = {
             "output_type": "chat",
             "input_type": "chat",
@@ -152,41 +128,20 @@ class LangflowAdapter(BaseAdapter):
                 logger.exception("Langflow API response decoding failed")
                 raise ValueError(f"Failed to decode JSON response: {e}") from e
 
-    async def ainvoke(
-        self,
-        input: List[AgentMessage],
-        session_id: Union[str, None],
-    ) -> AgentInvocationResponse:
-        """Asynchronous invocation method - invokes the Langflow agent and converts to List[AgentMessage]."""
-        logger.info(
-            "Langflow ainvoke called (session_id=%s, input_messages=%d)",
-            session_id,
-            len(input),
-        )
-        agent_input = self._prepare_input(input)
+    async def ainvoke(self, input: InvokeInput) -> InvokeOutput:
+        """Invoke Langflow agent and return response."""
+        logger.info("Langflow ainvoke called (session_id=%s)", input.session_id)
 
         try:
             agent_output = await self._aapi_call(
-                query=agent_input,
-                sessionId=session_id if session_id else uuid.uuid4(),
+                query=input.user_message_str(),
+                sessionId=input.session_id if input.session_id else str(uuid.uuid4()),
             )
-            logger.debug(
-                "Langflow API returned payload with top-level keys: %s",
-                list(agent_output.keys()),
-            )
-        except Exception as e:
-            logger.exception("Error calling Langflow endpoint")
-            raise RuntimeError(f"Error calling langflow endpoint: {e}") from e
-
-        try:
             agent_trajectory = self._convert_api_output_to_messages(agent_output)
             logger.info(
-                "Langflow conversion produced %d trajectory messages",
-                len(agent_trajectory),
-            )
-            return AgentInvocationResponse(
-                agent_trajectory=agent_trajectory,
+                "Langflow produced %d trajectory messages", len(agent_trajectory)
             )
+            return InvokeOutput(agent_trajectory=agent_trajectory)
         except Exception as e:
-            logger.exception("Error processing Langflow response")
-            raise RuntimeError(f"Error processing langflow response: {e}") from e
+            logger.exception("Error calling Langflow endpoint")
+            raise RuntimeError(f"Error calling Langflow endpoint: {e}") from e