Always include URI so that we can get by uri later
This commit is contained in:
parent
e70fa737b7
commit
cd2ff36684
1 changed files with 21 additions and 7 deletions
|
|
@ -4,7 +4,7 @@ from contextlib import asynccontextmanager
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import logfire
|
import logfire
|
||||||
from pydantic import BaseModel, TypeAdapter
|
from pydantic import BaseModel, Field, TypeAdapter
|
||||||
from pydantic_ai import Agent, RunContext
|
from pydantic_ai import Agent, RunContext
|
||||||
from pydantic_ai.messages import ModelMessage
|
from pydantic_ai.messages import ModelMessage
|
||||||
from pydantic_core import to_jsonable_python
|
from pydantic_core import to_jsonable_python
|
||||||
|
|
@ -12,7 +12,6 @@ from pydantic_core import to_jsonable_python
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.config import Config
|
from haiku.rag.config import Config
|
||||||
from haiku.rag.graph.common import get_model
|
from haiku.rag.graph.common import get_model
|
||||||
from haiku.rag.qa.agent import SearchResult
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
@ -40,6 +39,17 @@ logfire.instrument_pydantic_ai()
|
||||||
ModelMessagesTypeAdapter = TypeAdapter(list[ModelMessage])
|
ModelMessagesTypeAdapter = TypeAdapter(list[ModelMessage])
|
||||||
|
|
||||||
|
|
||||||
|
class SearchResult(BaseModel):
|
||||||
|
"""Search result with both title and URI for A2A agent."""
|
||||||
|
|
||||||
|
content: str = Field(description="The document text content")
|
||||||
|
score: float = Field(description="Relevance score (higher is more relevant)")
|
||||||
|
document_title: str | None = Field(
|
||||||
|
description="Human-readable document title", default=None
|
||||||
|
)
|
||||||
|
document_uri: str = Field(description="Document URI/path for get_full_document")
|
||||||
|
|
||||||
|
|
||||||
class AgentDependencies(BaseModel):
|
class AgentDependencies(BaseModel):
|
||||||
"""Dependencies for the A2A conversational agent."""
|
"""Dependencies for the A2A conversational agent."""
|
||||||
|
|
||||||
|
|
@ -70,15 +80,18 @@ Critical rules:
|
||||||
- Be concise and direct
|
- Be concise and direct
|
||||||
|
|
||||||
Citation Format:
|
Citation Format:
|
||||||
After your answer, include a "Sources:" section listing document URIs from search results.
|
After your answer, include a "Sources:" section listing documents from search results.
|
||||||
Format: "Sources:\n- [document_uri]"
|
Show both title and URI if available, otherwise just the URI.
|
||||||
|
Format: "Sources:\n- [document_title] ([document_uri])" or "Sources:\n- [document_uri]"
|
||||||
|
|
||||||
Example:
|
Example:
|
||||||
[Your answer here]
|
[Your answer here]
|
||||||
|
|
||||||
Sources:
|
Sources:
|
||||||
- /path/to/document.pdf
|
- Python Documentation (/guides/python.md)
|
||||||
- /another/document.md
|
- /guides/python-basics.md
|
||||||
|
|
||||||
|
Note: When using get_full_document, always use document_uri (not document_title).
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -195,7 +208,8 @@ def create_a2a_app(db_path: Path):
|
||||||
SearchResult(
|
SearchResult(
|
||||||
content=chunk.content,
|
content=chunk.content,
|
||||||
score=score,
|
score=score,
|
||||||
document_uri=(chunk.document_title or chunk.document_uri or ""),
|
document_title=chunk.document_title,
|
||||||
|
document_uri=(chunk.document_uri or ""),
|
||||||
)
|
)
|
||||||
for chunk, score in expanded_results
|
for chunk, score in expanded_results
|
||||||
]
|
]
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue