two version of R2R are here HEAD master

1 files changed, 554 insertions, 0 deletions

diff --git a/.venv/lib/python3.12/site-packages/sdk/sync_methods/retrieval.py b/.venv/lib/python3.12/site-packages/sdk/sync_methods/retrieval.py
new file mode 100644
index 00000000..4a927014
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/sdk/sync_methods/retrieval.py
@@ -0,0 +1,554 @@
+import json
+import uuid
+from typing import Any, Generator, Optional
+
+from shared.api.models import (
+    WrappedAgentResponse,
+    WrappedEmbeddingResponse,
+    WrappedLLMChatCompletion,
+    WrappedRAGResponse,
+    WrappedSearchResponse,
+)
+
+from ..models import (
+    AgentEvent,
+    CitationData,
+    CitationEvent,
+    Delta,
+    DeltaPayload,
+    FinalAnswerData,
+    FinalAnswerEvent,
+    GenerationConfig,
+    Message,
+    MessageData,
+    MessageDelta,
+    MessageEvent,
+    SearchMode,
+    SearchResultsData,
+    SearchResultsEvent,
+    SearchSettings,
+    ThinkingData,
+    ThinkingEvent,
+    ToolCallData,
+    ToolCallEvent,
+    ToolResultData,
+    ToolResultEvent,
+    UnknownEvent,
+)
+
+
+def parse_retrieval_event(raw: dict) -> Optional[AgentEvent]:
+    """
+    Convert a raw SSE event dict into a typed Pydantic model.
+
+    Example raw dict:
+        {
+          "event": "message",
+          "data": "{\"id\": \"msg_partial\", \"object\": \"agent.message.delta\", \"delta\": {...}}"
+        }
+    """
+    event_type = raw.get("event", "unknown")
+
+    # If event_type == "done", we usually return None to signal the SSE stream is finished.
+    if event_type == "done":
+        return None
+
+    # The SSE "data" is JSON-encoded, so parse it
+    data_str = raw.get("data", "")
+    try:
+        data_obj = json.loads(data_str)
+    except json.JSONDecodeError as e:
+        # You can decide whether to raise or return UnknownEvent
+        raise ValueError(f"Could not parse JSON in SSE event data: {e}") from e
+
+    # Now branch on event_type to build the right Pydantic model
+    if event_type == "search_results":
+        return SearchResultsEvent(
+            event=event_type,
+            data=SearchResultsData(**data_obj),
+        )
+    elif event_type == "message":
+        # Parse nested delta structure manually before creating MessageData
+        if "delta" in data_obj and isinstance(data_obj["delta"], dict):
+            delta_dict = data_obj["delta"]
+
+            # Convert content items to MessageDelta objects
+            if "content" in delta_dict and isinstance(
+                delta_dict["content"], list
+            ):
+                parsed_content = []
+                for item in delta_dict["content"]:
+                    if isinstance(item, dict):
+                        # Parse payload to DeltaPayload
+                        if "payload" in item and isinstance(
+                            item["payload"], dict
+                        ):
+                            payload_dict = item["payload"]
+                            item["payload"] = DeltaPayload(**payload_dict)
+                        parsed_content.append(MessageDelta(**item))
+
+                # Replace with parsed content
+                delta_dict["content"] = parsed_content
+
+            # Create properly typed Delta object
+            data_obj["delta"] = Delta(**delta_dict)
+
+        return MessageEvent(
+            event=event_type,
+            data=MessageData(**data_obj),
+        )
+    elif event_type == "citation":
+        return CitationEvent(event=event_type, data=CitationData(**data_obj))
+    elif event_type == "tool_call":
+        return ToolCallEvent(event=event_type, data=ToolCallData(**data_obj))
+    elif event_type == "tool_result":
+        return ToolResultEvent(
+            event=event_type, data=ToolResultData(**data_obj)
+        )
+    elif event_type == "thinking":
+        # Parse nested delta structure manually before creating ThinkingData
+        if "delta" in data_obj and isinstance(data_obj["delta"], dict):
+            delta_dict = data_obj["delta"]
+
+            # Convert content items to MessageDelta objects
+            if "content" in delta_dict and isinstance(
+                delta_dict["content"], list
+            ):
+                parsed_content = []
+                for item in delta_dict["content"]:
+                    if isinstance(item, dict):
+                        # Parse payload to DeltaPayload
+                        if "payload" in item and isinstance(
+                            item["payload"], dict
+                        ):
+                            payload_dict = item["payload"]
+                            item["payload"] = DeltaPayload(**payload_dict)
+                        parsed_content.append(MessageDelta(**item))
+
+                # Replace with parsed content
+                delta_dict["content"] = parsed_content
+
+            # Create properly typed Delta object
+            data_obj["delta"] = Delta(**delta_dict)
+
+        return ThinkingEvent(
+            event=event_type,
+            data=ThinkingData(**data_obj),
+        )
+    elif event_type == "final_answer":
+        return FinalAnswerEvent(
+            event=event_type, data=FinalAnswerData(**data_obj)
+        )
+    else:
+        # Fallback if it doesn't match any known event
+        return UnknownEvent(
+            event=event_type,
+            data=data_obj,
+        )
+
+
+def search_arg_parser(
+    query: str,
+    search_mode: Optional[str | SearchMode] = "custom",
+    search_settings: Optional[dict | SearchSettings] = None,
+) -> dict:
+    if search_mode and not isinstance(search_mode, str):
+        search_mode = search_mode.value
+
+    if search_settings and not isinstance(search_settings, dict):
+        search_settings = search_settings.model_dump()
+
+    data: dict[str, Any] = {
+        "query": query,
+        "search_settings": search_settings,
+    }
+    if search_mode:
+        data["search_mode"] = search_mode
+
+    return data
+
+
+def completion_arg_parser(
+    messages: list[dict | Message],
+    generation_config: Optional[dict | GenerationConfig] = None,
+) -> dict:
+    # FIXME: Needs a proper return type
+    cast_messages: list[Message] = [
+        Message(**msg) if isinstance(msg, dict) else msg for msg in messages
+    ]
+
+    if generation_config and not isinstance(generation_config, dict):
+        generation_config = generation_config.model_dump()
+
+    data: dict[str, Any] = {
+        "messages": [msg.model_dump() for msg in cast_messages],
+        "generation_config": generation_config,
+    }
+    return data
+
+
+def embedding_arg_parser(
+    text: str,
+) -> dict:
+    data: dict[str, Any] = {
+        "text": text,
+    }
+    return data
+
+
+def rag_arg_parser(
+    query: str,
+    rag_generation_config: Optional[dict | GenerationConfig] = None,
+    search_mode: Optional[str | SearchMode] = "custom",
+    search_settings: Optional[dict | SearchSettings] = None,
+    task_prompt: Optional[str] = None,
+    include_title_if_available: Optional[bool] = False,
+    include_web_search: Optional[bool] = False,
+) -> dict:
+    if rag_generation_config and not isinstance(rag_generation_config, dict):
+        rag_generation_config = rag_generation_config.model_dump()
+    if search_settings and not isinstance(search_settings, dict):
+        search_settings = search_settings.model_dump()
+
+    data: dict[str, Any] = {
+        "query": query,
+        "rag_generation_config": rag_generation_config,
+        "search_settings": search_settings,
+        "task_prompt": task_prompt,
+        "include_title_if_available": include_title_if_available,
+        "include_web_search": include_web_search,
+    }
+    if search_mode:
+        data["search_mode"] = search_mode
+    return data
+
+
+def agent_arg_parser(
+    message: Optional[dict | Message] = None,
+    rag_generation_config: Optional[dict | GenerationConfig] = None,
+    research_generation_config: Optional[dict | GenerationConfig] = None,
+    search_mode: Optional[str | SearchMode] = "custom",
+    search_settings: Optional[dict | SearchSettings] = None,
+    task_prompt: Optional[str] = None,
+    include_title_if_available: Optional[bool] = True,
+    conversation_id: Optional[str | uuid.UUID] = None,
+    max_tool_context_length: Optional[int] = None,
+    use_system_context: Optional[bool] = True,
+    rag_tools: Optional[list[str]] = None,
+    research_tools: Optional[list[str]] = None,
+    tools: Optional[list[str]] = None,  # For backward compatibility
+    mode: Optional[str] = "rag",
+    needs_initial_conversation_name: Optional[bool] = None,
+) -> dict:
+    if rag_generation_config and not isinstance(rag_generation_config, dict):
+        rag_generation_config = rag_generation_config.model_dump()
+    if research_generation_config and not isinstance(
+        research_generation_config, dict
+    ):
+        research_generation_config = research_generation_config.model_dump()
+    if search_settings and not isinstance(search_settings, dict):
+        search_settings = search_settings.model_dump()
+
+    data: dict[str, Any] = {
+        "rag_generation_config": rag_generation_config or {},
+        "search_settings": search_settings,
+        "task_prompt": task_prompt,
+        "include_title_if_available": include_title_if_available,
+        "conversation_id": (str(conversation_id) if conversation_id else None),
+        "max_tool_context_length": max_tool_context_length,
+        "use_system_context": use_system_context,
+        "mode": mode,
+    }
+
+    # Handle generation configs based on mode
+    if research_generation_config and mode == "research":
+        data["research_generation_config"] = research_generation_config
+
+    # Handle tool configurations
+    if rag_tools:
+        data["rag_tools"] = rag_tools
+    if research_tools:
+        data["research_tools"] = research_tools
+    if tools:  # Backward compatibility
+        data["tools"] = tools
+
+    if search_mode:
+        data["search_mode"] = search_mode
+
+    if needs_initial_conversation_name:
+        data["needs_initial_conversation_name"] = (
+            needs_initial_conversation_name
+        )
+
+    if message:
+        cast_message: Message = (
+            Message(**message) if isinstance(message, dict) else message
+        )
+        data["message"] = cast_message.model_dump()
+    return data
+
+
+class RetrievalSDK:
+    """SDK for interacting with documents in the v3 API."""
+
+    def __init__(self, client):
+        self.client = client
+
+    def search(
+        self,
+        query: str,
+        search_mode: Optional[str | SearchMode] = "custom",
+        search_settings: Optional[dict | SearchSettings] = None,
+    ) -> WrappedSearchResponse:
+        """Conduct a vector and/or graph search.
+
+        Args:
+            query (str): The query to search for.
+            search_settings (Optional[dict, SearchSettings]]): Vector search settings.
+
+        Returns:
+            WrappedSearchResponse
+        """
+
+        response_dict = self.client._make_request(
+            "POST",
+            "retrieval/search",
+            json=search_arg_parser(
+                query=query,
+                search_mode=search_mode,
+                search_settings=search_settings,
+            ),
+            version="v3",
+        )
+
+        return WrappedSearchResponse(**response_dict)
+
+    def completion(
+        self,
+        messages: list[dict | Message],
+        generation_config: Optional[dict | GenerationConfig] = None,
+    ) -> WrappedLLMChatCompletion:
+        cast_messages: list[Message] = [
+            Message(**msg) if isinstance(msg, dict) else msg
+            for msg in messages
+        ]
+
+        if generation_config and not isinstance(generation_config, dict):
+            generation_config = generation_config.model_dump()
+
+        data: dict[str, Any] = {
+            "messages": [msg.model_dump() for msg in cast_messages],
+            "generation_config": generation_config,
+        }
+        response_dict = self.client._make_request(
+            "POST",
+            "retrieval/completion",
+            json=completion_arg_parser(messages, generation_config),
+            version="v3",
+        )
+
+        return WrappedLLMChatCompletion(**response_dict)
+
+    def embedding(
+        self,
+        text: str,
+    ) -> WrappedEmbeddingResponse:
+        response_dict = self.client._make_request(
+            "POST",
+            "retrieval/embedding",
+            data=embedding_arg_parser(text),
+            version="v3",
+        )
+
+        return WrappedEmbeddingResponse(**response_dict)
+
+    def rag(
+        self,
+        query: str,
+        rag_generation_config: Optional[dict | GenerationConfig] = None,
+        search_mode: Optional[str | SearchMode] = "custom",
+        search_settings: Optional[dict | SearchSettings] = None,
+        task_prompt: Optional[str] = None,
+        include_title_if_available: Optional[bool] = False,
+        include_web_search: Optional[bool] = False,
+    ) -> (
+        WrappedRAGResponse
+        | Generator[
+            ThinkingEvent
+            | SearchResultsEvent
+            | MessageEvent
+            | CitationEvent
+            | FinalAnswerEvent
+            | ToolCallEvent
+            | ToolResultEvent
+            | UnknownEvent
+            | None,
+            None,
+            None,
+        ]
+    ):
+        """Conducts a Retrieval Augmented Generation (RAG) search with the
+        given query.
+
+        Args:
+            query (str): The query to search for.
+            rag_generation_config (Optional[dict | GenerationConfig]): RAG generation configuration.
+            search_settings (Optional[dict | SearchSettings]): Vector search settings.
+            task_prompt (Optional[str]): Task prompt override.
+            include_title_if_available (Optional[bool]): Include the title if available.
+
+        Returns:
+            WrappedRAGResponse | AsyncGenerator[RAGResponse, None]: The RAG response
+        """
+        data = rag_arg_parser(
+            query=query,
+            rag_generation_config=rag_generation_config,
+            search_mode=search_mode,
+            search_settings=search_settings,
+            task_prompt=task_prompt,
+            include_title_if_available=include_title_if_available,
+            include_web_search=include_web_search,
+        )
+        rag_generation_config = data.get("rag_generation_config")
+        if rag_generation_config and rag_generation_config.get(  # type: ignore
+            "stream", False
+        ):
+            raw_stream = self.client._make_streaming_request(
+                "POST",
+                "retrieval/rag",
+                json=data,
+                version="v3",
+            )
+            # Wrap the raw stream to parse each event
+            return (parse_retrieval_event(event) for event in raw_stream)
+
+        response_dict = self.client._make_request(
+            "POST",
+            "retrieval/rag",
+            json=data,
+            version="v3",
+        )
+
+        return WrappedRAGResponse(**response_dict)
+
+    def agent(
+        self,
+        message: Optional[dict | Message] = None,
+        rag_generation_config: Optional[dict | GenerationConfig] = None,
+        research_generation_config: Optional[dict | GenerationConfig] = None,
+        search_mode: Optional[str | SearchMode] = "custom",
+        search_settings: Optional[dict | SearchSettings] = None,
+        task_prompt: Optional[str] = None,
+        include_title_if_available: Optional[bool] = True,
+        conversation_id: Optional[str | uuid.UUID] = None,
+        max_tool_context_length: Optional[int] = None,
+        use_system_context: Optional[bool] = True,
+        # Tool configurations
+        rag_tools: Optional[list[str]] = None,
+        research_tools: Optional[list[str]] = None,
+        tools: Optional[list[str]] = None,  # For backward compatibility
+        mode: Optional[str] = "rag",
+        needs_initial_conversation_name: Optional[bool] = None,
+    ) -> (
+        WrappedAgentResponse
+        | Generator[
+            ThinkingEvent
+            | SearchResultsEvent
+            | MessageEvent
+            | CitationEvent
+            | FinalAnswerEvent
+            | ToolCallEvent
+            | ToolResultEvent
+            | UnknownEvent
+            | None,
+            None,
+            None,
+        ]
+    ):
+        """Performs a single turn in a conversation with a RAG agent.
+
+        Args:
+            message (Optional[dict | Message]): The message to send to the agent.
+            rag_generation_config (Optional[dict | GenerationConfig]): Configuration for RAG generation in 'rag' mode.
+            research_generation_config (Optional[dict | GenerationConfig]): Configuration for generation in 'research' mode.
+            search_mode (Optional[str | SearchMode]): Pre-configured search modes: "basic", "advanced", or "custom".
+            search_settings (Optional[dict | SearchSettings]): Vector search settings.
+            task_prompt (Optional[str]): Task prompt override.
+            include_title_if_available (Optional[bool]): Include the title if available.
+            conversation_id (Optional[str | uuid.UUID]): ID of the conversation for maintaining context.
+            max_tool_context_length (Optional[int]): Maximum context length for tool replies.
+            use_system_context (Optional[bool]): Whether to use system context in the prompt.
+            rag_tools (Optional[list[str]]): List of tools to enable for RAG mode.
+                Available tools: "search_file_knowledge", "content", "web_search", "web_scrape", "search_file_descriptions".
+            research_tools (Optional[list[str]]): List of tools to enable for Research mode.
+                Available tools: "rag", "reasoning", "critique", "python_executor".
+            tools (Optional[list[str]]): Deprecated. List of tools to execute.
+            mode (Optional[str]): Mode to use for generation: "rag" for standard retrieval or "research" for deep analysis.
+                Defaults to "rag".
+
+        Returns:
+            WrappedAgentResponse | AsyncGenerator[AgentEvent, None]: The agent response.
+        """
+        data = agent_arg_parser(
+            message=message,
+            rag_generation_config=rag_generation_config,
+            research_generation_config=research_generation_config,
+            search_mode=search_mode,
+            search_settings=search_settings,
+            task_prompt=task_prompt,
+            include_title_if_available=include_title_if_available,
+            conversation_id=conversation_id,
+            max_tool_context_length=max_tool_context_length,
+            use_system_context=use_system_context,
+            rag_tools=rag_tools,
+            research_tools=research_tools,
+            tools=tools,
+            mode=mode,
+            needs_initial_conversation_name=needs_initial_conversation_name,
+        )
+
+        # Determine if streaming is enabled
+        if search_mode:
+            data["search_mode"] = search_mode
+
+        if message:
+            cast_message: Message = (
+                Message(**message) if isinstance(message, dict) else message
+            )
+            data["message"] = cast_message.model_dump()
+
+        is_stream = False
+        if mode != "research":
+            if rag_generation_config:
+                if isinstance(rag_generation_config, dict):
+                    is_stream = rag_generation_config.get(  # type: ignore
+                        "stream", False
+                    )
+                else:
+                    is_stream = rag_generation_config.stream
+        else:
+            if research_generation_config:
+                if isinstance(research_generation_config, dict):
+                    is_stream = research_generation_config.get(  # type: ignore
+                        "stream", False
+                    )
+                else:
+                    is_stream = research_generation_config.stream
+
+        if is_stream:
+            raw_stream = self.client._make_streaming_request(
+                "POST",
+                "retrieval/agent",
+                json=data,
+                version="v3",
+            )
+            return (parse_retrieval_event(event) for event in raw_stream)
+
+        response_dict = self.client._make_request(
+            "POST",
+            "retrieval/agent",
+            json=data,
+            version="v3",
+        )
+
+        return WrappedAgentResponse(**response_dict)