501 lines
15 KiB
Python
501 lines
15 KiB
Python
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel
|
|
from pydantic_ai import RunContext
|
|
|
|
from haiku.rag.agents.research.models import Citation
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.store.models.chunk import SearchResult
|
|
from haiku.rag.tools.document import DocumentInfo
|
|
from haiku.rag.tools.qa import QAHistoryEntry
|
|
from haiku.skills.state import SkillRunDeps
|
|
|
|
|
|
class ResearchEntry(BaseModel):
|
|
question: str
|
|
title: str
|
|
executive_summary: str
|
|
|
|
|
|
class AnalysisEntry(BaseModel):
|
|
question: str
|
|
answer: str
|
|
program: str | None = None
|
|
|
|
|
|
async def find_relevant_prior_qa(
|
|
qa_history: list[QAHistoryEntry],
|
|
query: str,
|
|
config: AppConfig,
|
|
) -> list[QAHistoryEntry]:
|
|
from haiku.rag.embeddings import get_embedder
|
|
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD
|
|
from haiku.rag.utils import cosine_similarity
|
|
|
|
if not qa_history:
|
|
return []
|
|
|
|
embedder = get_embedder(config)
|
|
query_embedding = await embedder.embed_query(query)
|
|
|
|
to_embed = []
|
|
to_embed_indices = []
|
|
for i, qa in enumerate(qa_history):
|
|
if qa.question_embedding is None:
|
|
to_embed.append(qa.question)
|
|
to_embed_indices.append(i)
|
|
|
|
if to_embed:
|
|
new_embeddings = await embedder.embed_documents(to_embed)
|
|
for i, idx in enumerate(to_embed_indices):
|
|
qa_history[idx].question_embedding = new_embeddings[i]
|
|
|
|
matches = []
|
|
for qa in qa_history:
|
|
if qa.question_embedding is not None:
|
|
similarity = cosine_similarity(query_embedding, qa.question_embedding)
|
|
if similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD:
|
|
matches.append(qa)
|
|
|
|
return matches
|
|
|
|
|
|
async def skill_search(
|
|
db_path: Path,
|
|
config: AppConfig,
|
|
query: str,
|
|
limit: int | None = None,
|
|
document_filter: str | None = None,
|
|
) -> tuple[str, list[SearchResult]]:
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
|
results = await rag.search(
|
|
query,
|
|
limit=limit,
|
|
filter=document_filter,
|
|
)
|
|
results = await rag.expand_context(results)
|
|
|
|
formatted = "\n\n---\n\n".join(
|
|
r.format_for_agent(rank=i + 1, total=len(results))
|
|
for i, r in enumerate(results)
|
|
)
|
|
return formatted, list(results)
|
|
|
|
|
|
async def skill_list_documents(
|
|
db_path: Path,
|
|
config: AppConfig,
|
|
limit: int | None = None,
|
|
offset: int | None = None,
|
|
filter: str | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
|
documents = await rag.list_documents(limit, offset, filter=filter)
|
|
return [
|
|
{
|
|
"id": doc.id,
|
|
"title": doc.title,
|
|
"uri": doc.uri,
|
|
"metadata": doc.metadata,
|
|
"created_at": str(doc.created_at),
|
|
"updated_at": str(doc.updated_at),
|
|
}
|
|
for doc in documents
|
|
]
|
|
|
|
|
|
async def skill_get_document(
|
|
db_path: Path,
|
|
config: AppConfig,
|
|
query: str,
|
|
) -> dict[str, Any] | None:
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
|
document = await rag.resolve_document(query)
|
|
if document is None:
|
|
return None
|
|
return {
|
|
"id": document.id,
|
|
"content": document.content,
|
|
"title": document.title,
|
|
"uri": document.uri,
|
|
"metadata": document.metadata,
|
|
"created_at": str(document.created_at),
|
|
"updated_at": str(document.updated_at),
|
|
}
|
|
|
|
|
|
async def skill_ask(
|
|
db_path: Path,
|
|
config: AppConfig,
|
|
question: str,
|
|
qa_history: list[QAHistoryEntry] | None = None,
|
|
document_filter: str | None = None,
|
|
) -> tuple[str, list[Citation]]:
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.utils import format_citations
|
|
|
|
ask_question = question
|
|
if qa_history:
|
|
matches = await find_relevant_prior_qa(qa_history, question, config)
|
|
if matches:
|
|
prior_parts = []
|
|
for qa in matches:
|
|
part = f"Q: {qa.question}\nA: {qa.answer}"
|
|
if qa.citations:
|
|
part += "\n" + format_citations(qa.citations)
|
|
prior_parts.append(part)
|
|
ask_question = (
|
|
"Context from prior questions in this session:\n\n"
|
|
+ "\n\n---\n\n".join(prior_parts)
|
|
+ "\n\n---\n\nCurrent question: "
|
|
+ question
|
|
)
|
|
|
|
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
|
answer, citations = await rag.ask(
|
|
ask_question,
|
|
filter=document_filter,
|
|
)
|
|
|
|
return answer, citations
|
|
|
|
|
|
async def skill_research(
|
|
db_path: Path,
|
|
config: AppConfig,
|
|
question: str,
|
|
document_filter: str | None = None,
|
|
) -> tuple[str, str, str]:
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
|
report = await rag.research(question, filter=document_filter)
|
|
|
|
parts = [
|
|
f"# {report.title}",
|
|
f"\n## Executive Summary\n{report.executive_summary}",
|
|
]
|
|
if report.main_findings:
|
|
parts.append("\n## Main Findings")
|
|
for finding in report.main_findings:
|
|
parts.append(f"- {finding}")
|
|
if report.conclusions:
|
|
parts.append("\n## Conclusions")
|
|
for conclusion in report.conclusions:
|
|
parts.append(f"- {conclusion}")
|
|
if report.limitations:
|
|
parts.append("\n## Limitations")
|
|
for limitation in report.limitations:
|
|
parts.append(f"- {limitation}")
|
|
if report.recommendations:
|
|
parts.append("\n## Recommendations")
|
|
for rec in report.recommendations:
|
|
parts.append(f"- {rec}")
|
|
parts.append(f"\n## Sources\n{report.sources_summary}")
|
|
|
|
formatted = "\n".join(parts)
|
|
return formatted, report.title, report.executive_summary
|
|
|
|
|
|
async def skill_analyze(
|
|
db_path: Path,
|
|
config: AppConfig,
|
|
question: str,
|
|
document: str | None = None,
|
|
document_filter: str | None = None,
|
|
) -> tuple[str, str, str | None]:
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
|
documents = [document] if document else None
|
|
result = await rag.analyze(
|
|
question, documents=documents, filter=document_filter
|
|
)
|
|
output = result.answer
|
|
if result.program:
|
|
output += f"\n\nProgram:\n{result.program}"
|
|
|
|
return output, result.answer, result.program
|
|
|
|
|
|
def update_documents_state(
|
|
documents_state: list[DocumentInfo],
|
|
doc_dicts: list[dict[str, Any]],
|
|
) -> None:
|
|
for doc_dict in doc_dicts:
|
|
doc_info = DocumentInfo(
|
|
id=str(doc_dict["id"]),
|
|
title=doc_dict["title"] or "Untitled",
|
|
uri=doc_dict.get("uri") or "",
|
|
created=doc_dict.get("created_at", ""),
|
|
)
|
|
if not any(d.id == doc_info.id for d in documents_state):
|
|
documents_state.append(doc_info)
|
|
|
|
|
|
def _get_state(ctx: RunContext[SkillRunDeps], state_type: type[BaseModel]) -> Any:
|
|
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, state_type):
|
|
return ctx.deps.state
|
|
return None
|
|
|
|
|
|
def create_skill_extras(
|
|
db_path: Path,
|
|
config: AppConfig,
|
|
) -> dict[str, Any]:
|
|
"""Create non-tool utility functions bound to a specific database.
|
|
|
|
Returns a dict of values that can be attached to a Skill's extras:
|
|
|
|
Keys:
|
|
- 'db_path': path to the LanceDB used to configure the skill
|
|
- 'config': config passed to (or derived for) the skill
|
|
- 'list_documents': returns info for documents in the database
|
|
- 'visualize_chunk': returns visualizations for chunks in the database
|
|
"""
|
|
|
|
async def visualize_chunk(chunk_id: str) -> list:
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
|
chunk = await rag.get_chunk_by_id(chunk_id)
|
|
if chunk is None:
|
|
return []
|
|
return await rag.visualize_chunk(chunk)
|
|
|
|
async def list_documents(
|
|
limit: int | None = None,
|
|
offset: int | None = None,
|
|
filter: str | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
|
documents = await rag.list_documents(limit, offset, filter=filter)
|
|
return [
|
|
{
|
|
"id": doc.id,
|
|
"title": doc.title,
|
|
"uri": doc.uri,
|
|
"metadata": doc.metadata,
|
|
"created_at": str(doc.created_at),
|
|
"updated_at": str(doc.updated_at),
|
|
}
|
|
for doc in documents
|
|
]
|
|
|
|
return {
|
|
"db_path": db_path,
|
|
"config": config,
|
|
"visualize_chunk": visualize_chunk,
|
|
"list_documents": list_documents,
|
|
}
|
|
|
|
|
|
def create_skill_tools(
|
|
db_path: Path,
|
|
config: AppConfig,
|
|
state_type: type[BaseModel],
|
|
tool_names: list[str],
|
|
) -> dict[str, Any]:
|
|
"""Create tool closures for a skill.
|
|
|
|
Returns a dict mapping tool name to async callable.
|
|
Each tool extracts state from RunContext, calls the shared implementation,
|
|
and updates state.
|
|
"""
|
|
tools: dict[str, Any] = {}
|
|
|
|
if "search" in tool_names:
|
|
|
|
async def search(
|
|
ctx: RunContext[SkillRunDeps], query: str, limit: int | None = None
|
|
) -> str:
|
|
"""Search the knowledge base using hybrid search (vector + full-text).
|
|
|
|
Returns ranked results with content and metadata.
|
|
|
|
Args:
|
|
query: The search query.
|
|
limit: Maximum number of results.
|
|
"""
|
|
state = _get_state(ctx, state_type)
|
|
formatted, results = await skill_search(
|
|
db_path,
|
|
config,
|
|
query,
|
|
limit=limit,
|
|
document_filter=state.document_filter if state else None,
|
|
)
|
|
if state:
|
|
state.searches[query] = results
|
|
return formatted
|
|
|
|
tools["search"] = search
|
|
|
|
if "list_documents" in tool_names:
|
|
|
|
async def list_documents(
|
|
ctx: RunContext[SkillRunDeps],
|
|
limit: int | None = None,
|
|
offset: int | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""List documents in the knowledge base with optional pagination.
|
|
|
|
Args:
|
|
limit: Maximum number of documents to return.
|
|
offset: Number of documents to skip.
|
|
"""
|
|
state = _get_state(ctx, state_type)
|
|
result = await skill_list_documents(
|
|
db_path,
|
|
config,
|
|
limit,
|
|
offset,
|
|
filter=state.document_filter if state else None,
|
|
)
|
|
if state:
|
|
update_documents_state(state.documents, result)
|
|
return result
|
|
|
|
tools["list_documents"] = list_documents
|
|
|
|
if "get_document" in tool_names:
|
|
|
|
async def get_document(
|
|
ctx: RunContext[SkillRunDeps], query: str
|
|
) -> dict[str, Any] | None:
|
|
"""Retrieve a document by ID, title, or URI.
|
|
|
|
Args:
|
|
query: Document ID, title, or URI to look up.
|
|
"""
|
|
result = await skill_get_document(db_path, config, query)
|
|
if result is not None:
|
|
state = _get_state(ctx, state_type)
|
|
if state:
|
|
update_documents_state(state.documents, [result])
|
|
return result
|
|
|
|
tools["get_document"] = get_document
|
|
|
|
if "ask" in tool_names:
|
|
|
|
async def ask(ctx: RunContext[SkillRunDeps], question: str) -> str:
|
|
"""Ask a question and get an answer with citations from the knowledge base.
|
|
|
|
Args:
|
|
question: The question to ask.
|
|
"""
|
|
from haiku.rag.utils import format_citations
|
|
|
|
state = _get_state(ctx, state_type)
|
|
answer, citations = await skill_ask(
|
|
db_path,
|
|
config,
|
|
question,
|
|
qa_history=state.qa_history if state else None,
|
|
document_filter=state.document_filter if state else None,
|
|
)
|
|
|
|
if state:
|
|
next_index = len(state.citations) + 1
|
|
for citation in citations:
|
|
citation.index = next_index
|
|
next_index += 1
|
|
state.citations.extend(citations)
|
|
state.qa_history.append(
|
|
QAHistoryEntry(
|
|
question=question, answer=answer, citations=citations
|
|
)
|
|
)
|
|
|
|
if citations:
|
|
answer += "\n\n" + format_citations(citations)
|
|
|
|
return answer
|
|
|
|
tools["ask"] = ask
|
|
|
|
if "research" in tool_names:
|
|
|
|
async def research(ctx: RunContext[SkillRunDeps], question: str) -> str:
|
|
"""Conduct deep multi-agent research on a question.
|
|
|
|
Iteratively searches, analyzes, and synthesizes information from the
|
|
knowledge base to produce a comprehensive research report.
|
|
Only use when the user explicitly requests deep research.
|
|
|
|
Args:
|
|
question: The research question to investigate.
|
|
"""
|
|
state = _get_state(ctx, state_type)
|
|
formatted, title, executive_summary = await skill_research(
|
|
db_path,
|
|
config,
|
|
question,
|
|
document_filter=state.document_filter if state else None,
|
|
)
|
|
|
|
if state:
|
|
state.reports.append(
|
|
ResearchEntry(
|
|
question=question,
|
|
title=title,
|
|
executive_summary=executive_summary,
|
|
)
|
|
)
|
|
state.qa_history.append(
|
|
QAHistoryEntry(question=question, answer=executive_summary)
|
|
)
|
|
|
|
return formatted
|
|
|
|
tools["research"] = research
|
|
|
|
if "analyze" in tool_names:
|
|
|
|
async def analyze(
|
|
ctx: RunContext[SkillRunDeps],
|
|
question: str,
|
|
document: str | None = None,
|
|
) -> str:
|
|
"""Answer complex analytical questions using code execution.
|
|
|
|
Use this for questions requiring computation, aggregation, or
|
|
data traversal across documents.
|
|
|
|
Args:
|
|
question: The question to answer.
|
|
document: Optional document ID or title to pre-load for analysis.
|
|
"""
|
|
state = _get_state(ctx, state_type)
|
|
state_filter = state.document_filter if state else None
|
|
output, answer, program = await skill_analyze(
|
|
db_path,
|
|
config,
|
|
question,
|
|
document=document,
|
|
document_filter=state_filter,
|
|
)
|
|
if state:
|
|
state.analyses.append(
|
|
AnalysisEntry(
|
|
question=question,
|
|
answer=answer,
|
|
program=program,
|
|
)
|
|
)
|
|
|
|
return output
|
|
|
|
tools["analyze"] = analyze
|
|
|
|
return tools
|