109 lines
3.4 KiB
Python
109 lines
3.4 KiB
Python
import os
|
|
from functools import cache
|
|
from pathlib import Path
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
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.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
|
|
from haiku.skills.parser import parse_skill_md
|
|
|
|
AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
|
|
|
|
CRITICAL RULES:
|
|
1. For greetings or casual chat: respond directly WITHOUT using any tools
|
|
2. NEVER make up information - always use skills to get facts from the knowledge base
|
|
3. When a skill returns citations, always include them in your response
|
|
"""
|
|
|
|
_RAG_TOOLS = ["search", "list_documents", "get_document", "cite"]
|
|
|
|
|
|
def get_agent_preamble(config: AppConfig) -> str:
|
|
"""Build the main agent preamble, prepending domain_preamble if configured."""
|
|
if config.prompts.domain_preamble:
|
|
return f"{config.prompts.domain_preamble}\n\n{AGENT_PREAMBLE}"
|
|
return AGENT_PREAMBLE
|
|
|
|
|
|
class RAGState(BaseModel):
|
|
citation_index: dict[str, Citation] = Field(default_factory=dict)
|
|
citations: list[str] = Field(default_factory=list)
|
|
document_filter: str | None = None
|
|
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
|
|
|
|
|
|
STATE_TYPE = RAGState
|
|
STATE_NAMESPACE = "rag"
|
|
|
|
_skill_path = Path(__file__).parent / "rag"
|
|
|
|
|
|
@cache
|
|
def skill_metadata() -> SkillMetadata:
|
|
metadata, _ = parse_skill_md(_skill_path / "SKILL.md")
|
|
return metadata
|
|
|
|
|
|
@cache
|
|
def instructions() -> str | None:
|
|
_, instr = parse_skill_md(_skill_path / "SKILL.md")
|
|
return instr
|
|
|
|
|
|
def state_metadata() -> StateMetadata:
|
|
return StateMetadata(
|
|
namespace=STATE_NAMESPACE,
|
|
type=STATE_TYPE,
|
|
schema=STATE_TYPE.model_json_schema(),
|
|
)
|
|
|
|
|
|
def create_skill(
|
|
db_path: Path | None = None,
|
|
config: AppConfig | None = None,
|
|
) -> Skill:
|
|
"""Create a RAG skill for searching and analyzing documents.
|
|
|
|
Args:
|
|
db_path: Path to the LanceDB database. Resolved from:
|
|
1. This argument
|
|
2. HAIKU_RAG_DB environment variable
|
|
3. haiku.rag default (config.storage.data_dir / "haiku.rag.lancedb")
|
|
config: haiku.rag AppConfig instance. If None, uses get_config().
|
|
"""
|
|
from haiku.rag.config import get_config
|
|
from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan
|
|
from haiku.rag.skills._tools import create_skill_extras, create_skill_tools
|
|
|
|
if config is None:
|
|
config = get_config()
|
|
|
|
if db_path is None:
|
|
env_db = os.environ.get("HAIKU_RAG_DB")
|
|
if env_db:
|
|
db_path = Path(env_db).expanduser()
|
|
else:
|
|
db_path = config.storage.data_dir / "haiku.rag.lancedb"
|
|
|
|
tools = create_skill_tools(db_path, config, RAGState, _RAG_TOOLS)
|
|
extras = create_skill_extras(db_path, config)
|
|
|
|
skill_instructions = instructions()
|
|
if config.prompts.domain_preamble and skill_instructions:
|
|
skill_instructions = f"{config.prompts.domain_preamble}\n\n{skill_instructions}"
|
|
|
|
return Skill(
|
|
metadata=skill_metadata(),
|
|
source=SkillSource.ENTRYPOINT,
|
|
path=_skill_path,
|
|
instructions=skill_instructions,
|
|
tools=list(tools.values()),
|
|
extras=extras,
|
|
state_type=STATE_TYPE,
|
|
state_namespace=STATE_NAMESPACE,
|
|
deps_type=RAGRunDeps,
|
|
lifespan=make_rag_lifespan(db_path, config),
|
|
)
|