diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e23366821..f631625cd 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -12,7 +12,7 @@ repos: hooks: - id: check-typing-imports name: Check for Dict, List, or Union usage - entry: bash -c 'echo "Checking for typing imports..." && find . -name "*.py" | grep -v "/migrations/" | grep -v "/.venv/" | xargs grep -n "from typing.*import.*[^d]Dict\\|from typing.*import.*List" || exit 0 && echo "⚠️ Please import dict instead of Dict, list instead of List, and the logical OR operator" && exit 1' + entry: bash -c 'echo "Checking for typing imports..." && find . -name "*.py" | grep -v "/migrations/" | grep -v "/.venv/" | xargs grep -n "from typing.*import.*[^d]Dict\\|from typing.*import.*List\\|from typing.*import.*Union" || exit 0 && echo "⚠️ Please import dict instead of Dict, list instead of List, and the logical OR operator" && exit 1' language: system types: [python] pass_filenames: false diff --git a/py/sdk/asnyc_methods/retrieval.py b/py/sdk/asnyc_methods/retrieval.py index 7c393369d..d825a91f8 100644 --- a/py/sdk/asnyc_methods/retrieval.py +++ b/py/sdk/asnyc_methods/retrieval.py @@ -280,7 +280,7 @@ class RetrievalSDK: search_settings (Optional[dict | SearchSettings]): The search configuration object. task_prompt (Optional[str]): Optional custom prompt to override default. include_title_if_available (Optional[bool]): Include document titles from search results. - conversation_id (Optional[Union[str, uuid.UUID]]): ID of the conversation. + conversation_id (Optional[str | uuid.UUID]): ID of the conversation. tools (Optional[list[str]]): List of tools to execute (deprecated). rag_tools (Optional[list[str]]): List of tools to enable for RAG mode. research_tools (Optional[list[str]]): List of tools to enable for Research mode. diff --git a/py/sdk/sync_methods/retrieval.py b/py/sdk/sync_methods/retrieval.py index 1a0fdde14..b250dd70e 100644 --- a/py/sdk/sync_methods/retrieval.py +++ b/py/sdk/sync_methods/retrieval.py @@ -1,6 +1,6 @@ import json import uuid -from typing import Any, Generator, Optional, Union +from typing import Any, Generator, Optional from core.base.api.models import ( AgentEvent, @@ -231,7 +231,7 @@ def agent_arg_parser( search_settings: Optional[dict | SearchSettings] = None, task_prompt: Optional[str] = None, include_title_if_available: Optional[bool] = True, - conversation_id: Optional[Union[str, uuid.UUID]] = None, + 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, @@ -433,7 +433,7 @@ class RetrievalSDK: search_settings: Optional[dict | SearchSettings] = None, task_prompt: Optional[str] = None, include_title_if_available: Optional[bool] = True, - conversation_id: Optional[Union[str, uuid.UUID]] = None, + conversation_id: Optional[str | uuid.UUID] = None, max_tool_context_length: Optional[int] = None, use_system_context: Optional[bool] = True, # Tool configurations @@ -467,7 +467,7 @@ class RetrievalSDK: 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[Union[str, uuid.UUID]]): ID of the conversation for maintaining context. + 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. diff --git a/py/shared/api/models/retrieval/responses.py b/py/shared/api/models/retrieval/responses.py index 40892c31b..4253a15fe 100644 --- a/py/shared/api/models/retrieval/responses.py +++ b/py/shared/api/models/retrieval/responses.py @@ -1,4 +1,4 @@ -from typing import Any, Literal, Optional, Union +from typing import Any, Literal, Optional from pydantic import BaseModel, Field @@ -31,12 +31,10 @@ class Citation(R2RSerializable): description="The type of object, e.g. `citation`", ) payload: Optional[ - Union[ - ChunkSearchResult, - GraphSearchResult, - WebPageSearchResult, - ContextDocumentResult, - ] + ChunkSearchResult + | GraphSearchResult + | WebPageSearchResult + | ContextDocumentResult ] = Field( ..., description="The object payload and it's corresponding type" ) @@ -445,28 +443,27 @@ class ThinkingEvent(SSEEventBase): # Create a union type for all RAG events -RAGEvent = Union[ - SearchResultsEvent, - MessageEvent, - CitationEvent, - FinalAnswerEvent, - UnknownEvent, - ToolCallEvent, - ToolResultEvent, - ToolResultData, - ToolResultEvent, -] +RAGEvent = ( + SearchResultsEvent + | MessageEvent + | CitationEvent + | FinalAnswerEvent + | ToolCallEvent + | ToolResultEvent + | ToolResultData + | UnknownEvent +) -AgentEvent = Union[ - ThinkingEvent, - SearchResultsEvent, - MessageEvent, - CitationEvent, - FinalAnswerEvent, - ToolCallEvent, - ToolResultEvent, - UnknownEvent, -] +AgentEvent = ( + ThinkingEvent + | SearchResultsEvent + | MessageEvent + | CitationEvent + | FinalAnswerEvent + | ToolCallEvent + | ToolResultEvent + | UnknownEvent +) WrappedCompletionResponse = R2RResults[LLMChatCompletion] # Create wrapped versions of the responses