r2rpy/py/shared/api/models/retrieval/responses.py

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