two version of R2R are here HEAD master

1 files changed, 604 insertions, 0 deletions

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]]