diff options
Diffstat (limited to '.venv/lib/python3.12/site-packages/shared/api/models')
12 files changed, 1108 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/__init__.py b/.venv/lib/python3.12/site-packages/shared/api/models/__init__.py new file mode 100644 index 00000000..2d39dab1 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/__init__.py @@ -0,0 +1,194 @@ +from shared.api.models.auth.responses import ( + TokenResponse, + WrappedTokenResponse, +) +from shared.api.models.base import ( + GenericBooleanResponse, + GenericMessageResponse, + PaginatedR2RResult, + R2RResults, + WrappedBooleanResponse, + WrappedGenericMessageResponse, +) +from shared.api.models.graph.responses import ( + GraphResponse, + WrappedCommunitiesResponse, + WrappedCommunityResponse, + WrappedEntitiesResponse, + WrappedEntityResponse, + WrappedGraphResponse, + WrappedGraphsResponse, + WrappedRelationshipResponse, + WrappedRelationshipsResponse, +) +from shared.api.models.ingestion.responses import ( + IngestionResponse, + WrappedIngestionResponse, + WrappedMetadataUpdateResponse, + WrappedUpdateResponse, + WrappedVectorIndexResponse, + WrappedVectorIndicesResponse, +) +from shared.api.models.management.responses import ( + ChunkResponse, + CollectionResponse, + ConversationResponse, + MessageResponse, + PromptResponse, + ServerStats, + SettingsResponse, + WrappedAPIKeyResponse, + WrappedAPIKeysResponse, + WrappedChunkResponse, + WrappedChunksResponse, + WrappedCollectionResponse, + WrappedCollectionsResponse, + WrappedConversationMessagesResponse, + WrappedConversationResponse, + WrappedConversationsResponse, + WrappedDocumentResponse, + WrappedDocumentsResponse, + WrappedLimitsResponse, + WrappedLoginResponse, + WrappedMessageResponse, + WrappedPromptResponse, + WrappedPromptsResponse, + WrappedServerStatsResponse, + WrappedSettingsResponse, + WrappedUserResponse, + WrappedUsersResponse, +) +from shared.api.models.retrieval.responses import ( + AgentEvent, + AgentResponse, + AggregateSearchResult, + Citation, + CitationData, + CitationEvent, + Delta, + DeltaPayload, + FinalAnswerData, + FinalAnswerEvent, + MessageData, + MessageDelta, + MessageEvent, + RAGEvent, + RAGResponse, + SearchResultsData, + SearchResultsEvent, + SSEEventBase, + ThinkingData, + ThinkingEvent, + ToolCallData, + ToolCallEvent, + ToolResultData, + ToolResultEvent, + UnknownEvent, + WrappedAgentResponse, + WrappedDocumentSearchResponse, + WrappedEmbeddingResponse, + WrappedLLMChatCompletion, + WrappedRAGResponse, + WrappedSearchResponse, + WrappedVectorSearchResponse, +) + +__all__ = [ + # Generic Responses + "SSEEventBase", + "SearchResultsData", + "SearchResultsEvent", + "MessageDelta", + "MessageData", + "MessageEvent", + "DeltaPayload", + "Delta", + "CitationData", + "CitationEvent", + "FinalAnswerData", + "FinalAnswerEvent", + "ToolCallData", + "ToolCallEvent", + "ToolResultData", + "ToolResultEvent", + "ThinkingData", + "ThinkingEvent", + "AgentEvent", + "RAGEvent", + "UnknownEvent", + # Auth Responses + "GenericMessageResponse", + "TokenResponse", + "WrappedTokenResponse", + "WrappedGenericMessageResponse", + # Ingestion Responses + "IngestionResponse", + "WrappedIngestionResponse", + "WrappedUpdateResponse", + "WrappedVectorIndexResponse", + "WrappedVectorIndicesResponse", + "WrappedMetadataUpdateResponse", + "GraphResponse", + "WrappedGraphResponse", + "WrappedGraphsResponse", + "WrappedEntityResponse", + "WrappedEntitiesResponse", + "WrappedRelationshipResponse", + "WrappedRelationshipsResponse", + "WrappedCommunityResponse", + "WrappedCommunitiesResponse", + # Management Responses + "PromptResponse", + "ServerStats", + "SettingsResponse", + "ChunkResponse", + "CollectionResponse", + "ConversationResponse", + "MessageResponse", + "WrappedServerStatsResponse", + "WrappedSettingsResponse", + # Document Responses + "WrappedDocumentResponse", + "WrappedDocumentsResponse", + # Collection Responses + "WrappedCollectionResponse", + "WrappedCollectionsResponse", + # Prompt Responses + "WrappedPromptResponse", + "WrappedPromptsResponse", + # Chunk Responses + "WrappedChunkResponse", + "WrappedChunksResponse", + # Conversation Responses + "WrappedConversationMessagesResponse", + "WrappedConversationResponse", + "WrappedConversationsResponse", + # User Responses + "WrappedUserResponse", + "WrappedAPIKeyResponse", + "WrappedLimitsResponse", + "WrappedAPIKeysResponse", + "WrappedLoginResponse", + "WrappedUsersResponse", + "WrappedMessageResponse", + # Base Responses + "PaginatedR2RResult", + "R2RResults", + "GenericBooleanResponse", + "GenericMessageResponse", + "WrappedBooleanResponse", + "WrappedGenericMessageResponse", + # TODO: Clean up the following responses + # Retrieval Responses + "RAGResponse", + "Citation", + "WrappedRAGResponse", + "AgentResponse", + "AggregateSearchResult", + "WrappedSearchResponse", + "WrappedDocumentSearchResponse", + "WrappedVectorSearchResponse", + "WrappedAgentResponse", + "WrappedLLMChatCompletion", + "WrappedEmbeddingResponse", +] diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/auth/__init__.py b/.venv/lib/python3.12/site-packages/shared/api/models/auth/__init__.py new file mode 100644 index 00000000..e69de29b --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/auth/__init__.py diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/auth/responses.py b/.venv/lib/python3.12/site-packages/shared/api/models/auth/responses.py new file mode 100644 index 00000000..2d448945 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/auth/responses.py @@ -0,0 +1,13 @@ +from pydantic import BaseModel + +from shared.abstractions import Token +from shared.api.models.base import R2RResults + + +class TokenResponse(BaseModel): + access_token: Token + refresh_token: Token + + +# Create wrapped versions of each response +WrappedTokenResponse = R2RResults[TokenResponse] diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/base.py b/.venv/lib/python3.12/site-packages/shared/api/models/base.py new file mode 100644 index 00000000..e0493d0b --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/base.py @@ -0,0 +1,26 @@ +from typing import Generic, TypeVar + +from pydantic import BaseModel + +T = TypeVar("T") + + +class R2RResults(BaseModel, Generic[T]): + results: T + + +class PaginatedR2RResult(BaseModel, Generic[T]): + results: T + total_entries: int + + +class GenericBooleanResponse(BaseModel): + success: bool + + +class GenericMessageResponse(BaseModel): + message: str + + +WrappedBooleanResponse = R2RResults[GenericBooleanResponse] +WrappedGenericMessageResponse = R2RResults[GenericMessageResponse] diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/graph/__init__.py b/.venv/lib/python3.12/site-packages/shared/api/models/graph/__init__.py new file mode 100644 index 00000000..e69de29b --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/graph/__init__.py diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/graph/responses.py b/.venv/lib/python3.12/site-packages/shared/api/models/graph/responses.py new file mode 100644 index 00000000..a2272833 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/graph/responses.py @@ -0,0 +1,31 @@ +from datetime import datetime +from typing import Optional +from uuid import UUID + +from pydantic import BaseModel + +from shared.abstractions.graph import Community, Entity, Relationship +from shared.api.models.base import PaginatedR2RResult, R2RResults + +WrappedEntityResponse = R2RResults[Entity] +WrappedEntitiesResponse = PaginatedR2RResult[list[Entity]] +WrappedRelationshipResponse = R2RResults[Relationship] +WrappedRelationshipsResponse = PaginatedR2RResult[list[Relationship]] +WrappedCommunityResponse = R2RResults[Community] +WrappedCommunitiesResponse = PaginatedR2RResult[list[Community]] + + +class GraphResponse(BaseModel): + id: UUID + collection_id: UUID + name: str + description: Optional[str] + status: str + created_at: datetime + updated_at: datetime + document_ids: list[UUID] + + +# Graph Responses +WrappedGraphResponse = R2RResults[GraphResponse] +WrappedGraphsResponse = PaginatedR2RResult[list[GraphResponse]] diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/ingestion/__init__.py b/.venv/lib/python3.12/site-packages/shared/api/models/ingestion/__init__.py new file mode 100644 index 00000000..e69de29b --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/ingestion/__init__.py diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/ingestion/responses.py b/.venv/lib/python3.12/site-packages/shared/api/models/ingestion/responses.py new file mode 100644 index 00000000..091e48e7 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/ingestion/responses.py @@ -0,0 +1,72 @@ +from typing import Any, Optional, TypeVar +from uuid import UUID + +from pydantic import BaseModel, Field + +from shared.api.models.base import PaginatedR2RResult, R2RResults + +T = TypeVar("T") + + +class IngestionResponse(BaseModel): + message: str = Field( + ..., + description="A message describing the result of the ingestion request.", + ) + task_id: Optional[UUID] = Field( + None, + description="The task ID of the ingestion request.", + ) + document_id: UUID = Field( + ..., + description="The ID of the document that was ingested.", + ) + + class Config: + json_schema_extra = { + "example": { + "message": "Ingestion task queued successfully.", + "task_id": "c68dc72e-fc23-5452-8f49-d7bd46088a96", + "document_id": "9fbe403b-c11c-5aae-8ade-ef22980c3ad1", + } + } + + +class UpdateResponse(BaseModel): + message: str = Field( + ..., + description="A message describing the result of the ingestion request.", + ) + task_id: Optional[UUID] = Field( + None, + description="The task ID of the ingestion request.", + ) + document_ids: list[UUID] = Field( + ..., + description="The ID of the document that was ingested.", + ) + + class Config: + json_schema_extra = { + "example": { + "message": "Update task queued successfully.", + "task_id": "c68dc72e-fc23-5452-8f49-d7bd46088a96", + "document_ids": ["9fbe403b-c11c-5aae-8ade-ef22980c3ad1"], + } + } + + +class VectorIndexResponse(BaseModel): + index: dict[str, Any] + + +class VectorIndicesResponse(BaseModel): + indices: list[VectorIndexResponse] + + +WrappedIngestionResponse = R2RResults[IngestionResponse] +WrappedMetadataUpdateResponse = R2RResults[IngestionResponse] +WrappedUpdateResponse = R2RResults[UpdateResponse] + +WrappedVectorIndexResponse = R2RResults[VectorIndexResponse] +WrappedVectorIndicesResponse = PaginatedR2RResult[VectorIndicesResponse] diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/management/__init__.py b/.venv/lib/python3.12/site-packages/shared/api/models/management/__init__.py new file mode 100644 index 00000000..e69de29b --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/management/__init__.py diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/management/responses.py b/.venv/lib/python3.12/site-packages/shared/api/models/management/responses.py new file mode 100644 index 00000000..5e8b67a2 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/management/responses.py @@ -0,0 +1,168 @@ +from datetime import datetime +from typing import Any, Optional +from uuid import UUID + +from pydantic import BaseModel + +from shared.abstractions.document import DocumentResponse +from shared.abstractions.llm import Message +from shared.abstractions.user import Token, User +from shared.api.models.base import PaginatedR2RResult, R2RResults + + +class PromptResponse(BaseModel): + id: UUID + name: str + template: str + created_at: datetime + updated_at: datetime + input_types: dict[str, str] + + +class ServerStats(BaseModel): + start_time: datetime + uptime_seconds: float + cpu_usage: float + memory_usage: float + + +class SettingsResponse(BaseModel): + config: dict[str, Any] + prompts: dict[str, Any] + r2r_project_name: str + # r2r_version: str + + +class ChunkResponse(BaseModel): + id: UUID + document_id: UUID + owner_id: UUID + collection_ids: list[UUID] + text: str + metadata: dict[str, Any] + vector: Optional[list[float]] = None + + +class CollectionResponse(BaseModel): + id: UUID + owner_id: Optional[UUID] + name: str + description: Optional[str] + graph_cluster_status: str + graph_sync_status: str + created_at: datetime + updated_at: datetime + user_count: int + document_count: int + + +class ConversationResponse(BaseModel): + id: UUID + created_at: datetime + user_id: Optional[UUID] = None + name: Optional[str] = None + + +class MessageResponse(BaseModel): + id: UUID + message: Message + metadata: dict[str, Any] = {} + + +class ApiKey(BaseModel): + public_key: str + api_key: str + key_id: str + name: Optional[str] = None + + +class ApiKeyNoPriv(BaseModel): + public_key: str + key_id: str + name: Optional[str] = None + updated_at: datetime + description: Optional[str] = None + + +class LoginResponse(BaseModel): + access_token: Token + refresh_token: Token + + +class UsageLimit(BaseModel): + used: int + limit: int + remaining: int + + +class StorageTypeLimit(BaseModel): + limit: int + used: int + remaining: int + + +class StorageLimits(BaseModel): + chunks: StorageTypeLimit + documents: StorageTypeLimit + collections: StorageTypeLimit + + +class RouteUsage(BaseModel): + route_per_min: UsageLimit + monthly_limit: UsageLimit + + +class Usage(BaseModel): + global_per_min: UsageLimit + monthly_limit: UsageLimit + routes: dict[str, RouteUsage] + + +class SystemDefaults(BaseModel): + global_per_min: int + route_per_min: Optional[int] + monthly_limit: int + + +class LimitsResponse(BaseModel): + storage_limits: StorageLimits + system_defaults: SystemDefaults + user_overrides: dict + effective_limits: SystemDefaults + usage: Usage + + +# Chunk Responses +WrappedChunkResponse = R2RResults[ChunkResponse] +WrappedChunksResponse = PaginatedR2RResult[list[ChunkResponse]] + +# Collection Responses +WrappedCollectionResponse = R2RResults[CollectionResponse] +WrappedCollectionsResponse = PaginatedR2RResult[list[CollectionResponse]] + +# Conversation Responses +WrappedConversationMessagesResponse = R2RResults[list[MessageResponse]] +WrappedConversationResponse = R2RResults[ConversationResponse] +WrappedConversationsResponse = PaginatedR2RResult[list[ConversationResponse]] +WrappedMessageResponse = R2RResults[MessageResponse] +WrappedMessagesResponse = PaginatedR2RResult[list[MessageResponse]] + +# Document Responses +WrappedDocumentResponse = R2RResults[DocumentResponse] +WrappedDocumentsResponse = PaginatedR2RResult[list[DocumentResponse]] + +# Prompt Responses +WrappedPromptResponse = R2RResults[PromptResponse] +WrappedPromptsResponse = PaginatedR2RResult[list[PromptResponse]] + +# System Responses +WrappedSettingsResponse = R2RResults[SettingsResponse] +WrappedServerStatsResponse = R2RResults[ServerStats] + +# User Responses +WrappedUserResponse = R2RResults[User] +WrappedUsersResponse = PaginatedR2RResult[list[User]] +WrappedAPIKeyResponse = R2RResults[ApiKey] +WrappedAPIKeysResponse = PaginatedR2RResult[list[ApiKeyNoPriv]] +WrappedLoginResponse = R2RResults[LoginResponse] +WrappedLimitsResponse = R2RResults[LimitsResponse] diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/retrieval/__init__.py b/.venv/lib/python3.12/site-packages/shared/api/models/retrieval/__init__.py new file mode 100644 index 00000000..e69de29b --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/retrieval/__init__.py diff --git a/.venv/lib/python3.12/site-packages/shared/api/models/retrieval/responses.py b/.venv/lib/python3.12/site-packages/shared/api/models/retrieval/responses.py new file mode 100644 index 00000000..f695ebfb --- /dev/null +++ b/.venv/lib/python3.12/site-packages/shared/api/models/retrieval/responses.py @@ -0,0 +1,604 @@ +from typing import Any, Literal, Optional + +from pydantic import BaseModel, Field + +from shared.abstractions import ( + AggregateSearchResult, + ChunkSearchResult, + GraphSearchResult, + LLMChatCompletion, + Message, + WebPageSearchResult, +) +from shared.api.models.base import R2RResults +from shared.api.models.management.responses import DocumentResponse + +from ....abstractions import R2RSerializable + + +class CitationSpan(R2RSerializable): + """Represents a single occurrence of a citation in text.""" + + start_index: int = Field( + ..., description="Starting character index of the citation" + ) + end_index: int = Field( + ..., description="Ending character index of the citation" + ) + context_start: int = Field( + ..., description="Starting index of the surrounding context" + ) + context_end: int = Field( + ..., description="Ending index of the surrounding context" + ) + + +class Citation(R2RSerializable): + """ + Represents a citation reference in the RAG response. + + The first time a citation appears, it includes the full payload. + Subsequent appearances only include the citation ID and span information. + """ + + # Basic identification + id: str = Field( + ..., description="The short ID of the citation (e.g., 'e41ac2d')" + ) + object: str = Field( + "citation", description="The type of object, always 'citation'" + ) + + # Optimize payload delivery + is_new: bool = Field( + True, + description="Whether this is the first occurrence of this citation", + ) + + # Position information + span: Optional[CitationSpan] = Field( + None, description="Position of this citation occurrence in the text" + ) + + # Source information - only included for first occurrence + source_type: Optional[str] = Field( + None, description="Type of source: 'chunk', 'graph', 'web', or 'doc'" + ) + + # Full payload - only included for first occurrence + payload: ( + ChunkSearchResult + | GraphSearchResult + | WebPageSearchResult + | DocumentResponse + | dict[str, Any] + | None + ) = Field( + None, + description="The complete source object (only included for new citations)", + ) + + class Config: + extra = "ignore" + json_schema_extra = { + "example": { + "id": "e41ac2d", + "object": "citation", + "is_new": True, + "span": { + "start_index": 120, + "end_index": 129, + "context_start": 80, + "context_end": 180, + }, + "source_type": "chunk", + "payload": { + "id": "e41ac2d1-full-id", + "text": "The study found significant improvements...", + "metadata": {"title": "Research Paper"}, + }, + } + } + + +# class Citation(R2RSerializable): +# """Represents a single citation reference in the RAG response. + +# Combines both bracket metadata (start/end offsets, snippet range) and the +# mapped source fields (id, doc ID, chunk text, etc.). +# """ + +# # Bracket references +# id: str = Field(..., description="The ID of the citation object") +# object: str = Field( +# ..., +# description="The type of object, e.g. `citation`", +# ) +# payload: ( +# ChunkSearchResult +# | GraphSearchResult +# | WebPageSearchResult +# | DocumentResponse +# | None +# ) = Field( +# ..., description="The object payload and it's corresponding type" +# ) + +# class Config: +# extra = "ignore" # This tells Pydantic to ignore extra fields +# json_schema_extra = { +# "example": { +# "id": "cit.abcd123", +# "object": "citation", +# "payload": "ChunkSearchResult(...)", +# } +# } + + +class RAGResponse(R2RSerializable): + generated_answer: str = Field( + ..., description="The generated completion from the RAG process" + ) + search_results: AggregateSearchResult = Field( + ..., description="The search results used for the RAG process" + ) + citations: Optional[list[Citation]] = Field( + None, + description="Structured citation metadata, if you do citation extraction.", + ) + metadata: dict = Field( + default_factory=dict, + description="Additional data returned by the LLM provider", + ) + completion: str = Field( + ..., + description="The generated completion from the RAG process", + # deprecated=True, + ) + + class Config: + json_schema_extra = { + "example": { + "generated_answer": "The capital of France is Paris.", + "search_results": { + "chunk_search_results": [ + { + "index": 1, + "start_index": 25, + "end_index": 28, + "uri": "https://example.com/doc1", + "title": "example_document_1.pdf", + "license": "CC-BY-4.0", + } + ], + "graph_search_results": [ + { + "content": { + "id": "3f3d47f3-8baf-58eb-8bc2-0171fb1c6e09", + "name": "Entity Name", + "description": "Entity Description", + "metadata": {}, + }, + "result_type": "entity", + "chunk_ids": [ + "c68dc72e-fc23-5452-8f49-d7bd46088a96" + ], + "metadata": { + "associated_query": "What is the capital of France?" + }, + } + ], + "web_search_results": [ + { + "title": "Page Title", + "link": "https://example.com/page", + "snippet": "Page snippet", + "position": 1, + "date": "2021-01-01", + "sitelinks": [ + { + "title": "Sitelink Title", + "link": "https://example.com/sitelink", + } + ], + } + ], + "document_search_results": [ + { + "document": { + "id": "3f3d47f3-8baf-58eb-8bc2-0171fb1c6e09", + "title": "Document Title", + "chunks": ["Chunk 1", "Chunk 2"], + "metadata": {}, + }, + } + ], + }, + "citations": [ + { + "index": 1, + "rawIndex": 9, + "startIndex": 393, + "endIndex": 396, + "snippetStartIndex": 320, + "snippetEndIndex": 418, + "sourceType": "chunk", + "id": "e760bb76-1c6e-52eb-910d-0ce5b567011b", + "document_id": "e43864f5-a36f-548e-aacd-6f8d48b30c7f", + "owner_id": "2acb499e-8428-543b-bd85-0d9098718220", + "collection_ids": [ + "122fdf6a-e116-546b-a8f6-e4cb2e2c0a09" + ], + "score": 0.64, + "text": "Document Title: DeepSeek_R1.pdf\n\nText: could achieve an accuracy of ...", + "metadata": { + "title": "DeepSeek_R1.pdf", + "license": "CC-BY-4.0", + "chunk_order": 68, + "document_type": "pdf", + }, + } + ], + "metadata": { + "id": "chatcmpl-example123", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": {"role": "assistant"}, + } + ], + }, + "completion": "TO BE DEPRECATED", + } + } + + +class AgentResponse(R2RSerializable): + messages: list[Message] = Field(..., description="Agent response messages") + conversation_id: str = Field( + ..., description="The conversation ID for the RAG agent response" + ) + + class Config: + json_schema_extra = { + "example": { + "messages": [ + { + "role": "assistant", + "content": """Aristotle (384–322 BC) was an Ancient + Greek philosopher and polymath whose contributions + have had a profound impact on various fields of + knowledge. + Here are some key points about his life and work: + \n\n1. **Early Life**: Aristotle was born in 384 BC in + Stagira, Chalcidice, which is near modern-day + Thessaloniki, Greece. His father, Nicomachus, was the + personal physician to King Amyntas of Macedon, which + exposed Aristotle to medical and biological knowledge + from a young age [C].\n\n2. **Education and Career**: + After the death of his parents, Aristotle was sent to + Athens to study at Plato's Academy, where he remained + for about 20 years. After Plato's death, Aristotle + left Athens and eventually became the tutor of + Alexander the Great [C]. + \n\n3. **Philosophical Contributions**: Aristotle + founded the Lyceum in Athens, where he established the + Peripatetic school of philosophy. His works cover a + wide range of subjects, including metaphysics, ethics, + politics, logic, biology, and aesthetics. His writings + laid the groundwork for many modern scientific and + philosophical inquiries [A].\n\n4. **Legacy**: + Aristotle's influence extends beyond philosophy to the + natural sciences, linguistics, economics, and + psychology. His method of systematic observation and + analysis has been foundational to the development of + modern science [A].\n\nAristotle's comprehensive + approach to knowledge and his systematic methodology + have earned him a lasting legacy as one of the + greatest philosophers of all time.\n\nSources: + \n- [A] Aristotle's broad range of writings and + influence on modern science.\n- [C] Details about + Aristotle's early life and education.""", + "name": None, + "function_call": None, + "tool_calls": None, + "metadata": { + "citations": [ + { + "index": 1, + "rawIndex": 9, + "startIndex": 393, + "endIndex": 396, + "snippetStartIndex": 320, + "snippetEndIndex": 418, + "sourceType": "chunk", + "id": "e760bb76-1c6e-52eb-910d-0ce5b567011b", + "document_id": """ + e43864f5-a36f-548e-aacd-6f8d48b30c7f + """, + "owner_id": """ + 2acb499e-8428-543b-bd85-0d9098718220 + """, + "collection_ids": [ + "122fdf6a-e116-546b-a8f6-e4cb2e2c0a09" + ], + "score": 0.64, + "text": """ + Document Title: DeepSeek_R1.pdf + \n\nText: could achieve an accuracy of ... + """, + "metadata": { + "title": "DeepSeek_R1.pdf", + "license": "CC-BY-4.0", + "chunk_order": 68, + "document_type": "pdf", + }, + } + ], + "aggregated_search_results": { + "chunk_search_results": [ + { + "id": "3f3d47f3-8baf-58eb-8bc2-0171fb1c6e09", + "document_id": "3e157b3a-8469-51db-90d9-52e7d896b49b", + "owner_id": "2acb499e-8428-543b-bd85-0d9098718220", + "collection_ids": [], + "score": 0.23943702876567796, + "text": "Example text from the document", + "metadata": { + "title": "example_document.pdf", + "associated_query": "What is the capital of France?", + }, + } + ], + "graph_search_results": [ + { + "content": { + "id": "3f3d47f3-8baf-58eb-8bc2-0171fb1c6e09", + "name": "Entity Name", + "description": "Entity Description", + "metadata": {}, + }, + "result_type": "entity", + "chunk_ids": [ + "c68dc72e-fc23-5452-8f49-d7bd46088a96" + ], + "metadata": { + "associated_query": "What is the capital of France?" + }, + } + ], + "web_search_results": [ + { + "title": "Page Title", + "link": "https://example.com/page", + "snippet": "Page snippet", + "position": 1, + "date": "2021-01-01", + "sitelinks": [ + { + "title": "Sitelink Title", + "link": "https://example.com/sitelink", + } + ], + } + ], + "document_search_results": [ + { + "document": { + "id": "3f3d47f3-8baf-58eb-8bc2-0171fb1c6e09", + "title": "Document Title", + "chunks": ["Chunk 1", "Chunk 2"], + "metadata": {}, + }, + } + ], + }, + }, + }, + ], + "conversation_id": "a32b4c5d-6e7f-8a9b-0c1d-2e3f4a5b6c7d", + } + } + + +class DocumentSearchResult(BaseModel): + document_id: str = Field( + ..., + description="The document ID", + ) + metadata: Optional[dict] = Field( + None, + description="The metadata of the document", + ) + score: float = Field( + ..., + description="The score of the document", + ) + + +# A generic base model for SSE events +class SSEEventBase(BaseModel): + event: str + data: Any + + +# Model for the search results event +class SearchResultsData(BaseModel): + id: str + object: str + data: AggregateSearchResult + + +class SearchResultsEvent(SSEEventBase): + event: Literal["search_results"] + data: SearchResultsData + + +class DeltaPayload(BaseModel): + value: str + annotations: list[Any] + + +# Model for message events (partial tokens) +class MessageDelta(BaseModel): + type: str + payload: DeltaPayload + + +class Delta(BaseModel): + content: list[MessageDelta] + + +class MessageData(BaseModel): + id: str + object: str + delta: Delta + + +class MessageEvent(SSEEventBase): + event: Literal["message"] + data: MessageData + + +# Update CitationSpan model for SSE events +class CitationSpanData(BaseModel): + start: int = Field( + ..., description="Starting character index of the citation" + ) + end: int = Field(..., description="Ending character index of the citation") + context_start: Optional[int] = Field( + None, description="Starting index of surrounding context" + ) + context_end: Optional[int] = Field( + None, description="Ending index of surrounding context" + ) + + +# Update CitationData model +class CitationData(BaseModel): + id: str = Field( + ..., description="The short ID of the citation (e.g., 'e41ac2d')" + ) + object: str = Field( + "citation", description="The type of object, always 'citation'" + ) + + # New fields from the enhanced Citation model + is_new: Optional[bool] = Field( + None, + description="Whether this is the first occurrence of this citation", + ) + + span: Optional[CitationSpanData] = Field( + None, description="Position of this citation occurrence in the text" + ) + + source_type: Optional[str] = Field( + None, description="Type of source: 'chunk', 'graph', 'web', or 'doc'" + ) + + # Optional payload field, only for first occurrence + payload: Optional[Any] = Field( + None, + description="The complete source object (only included for new citations)", + ) + + # For backward compatibility, maintain the existing fields + class Config: + populate_by_name = True + extra = "ignore" + + +# CitationEvent remains the same, but now using the updated CitationData +class CitationEvent(SSEEventBase): + event: Literal["citation"] + data: CitationData + + +# Model for the final answer event +class FinalAnswerData(BaseModel): + generated_answer: str + citations: list[Citation] # refine if you have a citation model + + +class FinalAnswerEvent(SSEEventBase): + event: Literal["final_answer"] + data: FinalAnswerData + + +# "tool_call" event +class ToolCallData(BaseModel): + tool_call_id: str + name: str + arguments: Any # If JSON arguments, use dict[str, Any], or str if needed + + +class ToolCallEvent(SSEEventBase): + event: Literal["tool_call"] + data: ToolCallData + + +# "tool_result" event +class ToolResultData(BaseModel): + tool_call_id: str + role: Literal["tool", "function"] + content: str + + +class ToolResultEvent(SSEEventBase): + event: Literal["tool_result"] + data: ToolResultData + + +# Optionally, define a fallback model for unrecognized events +class UnknownEvent(SSEEventBase): + pass + + +# 1) Define a new ThinkingEvent type +class ThinkingData(BaseModel): + id: str + object: str + delta: Delta + + +class ThinkingEvent(SSEEventBase): + event: str = "thinking" + data: ThinkingData + + +# Create a union type for all RAG events +RAGEvent = ( + SearchResultsEvent + | MessageEvent + | CitationEvent + | FinalAnswerEvent + | UnknownEvent + | ToolCallEvent + | ToolResultEvent + | ToolResultData + | ToolResultEvent +) + +AgentEvent = ( + ThinkingEvent + | SearchResultsEvent + | MessageEvent + | CitationEvent + | FinalAnswerEvent + | ToolCallEvent + | ToolResultEvent + | UnknownEvent +) + +WrappedCompletionResponse = R2RResults[LLMChatCompletion] +# Create wrapped versions of the responses +WrappedVectorSearchResponse = R2RResults[list[ChunkSearchResult]] +WrappedSearchResponse = R2RResults[AggregateSearchResult] +# FIXME: This is returning DocumentResponse, but should be DocumentSearchResult +WrappedDocumentSearchResponse = R2RResults[list[DocumentResponse]] +WrappedRAGResponse = R2RResults[RAGResponse] +WrappedAgentResponse = R2RResults[AgentResponse] +WrappedLLMChatCompletion = R2RResults[LLMChatCompletion] +WrappedEmbeddingResponse = R2RResults[list[float]] |