haiku.rag 0.9.3__py3-none-any.whl → 0.10.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of haiku.rag might be problematic. Click here for more details.
- haiku/rag/app.py +64 -18
- haiku/rag/cli.py +67 -30
- haiku/rag/client.py +63 -21
- haiku/rag/config.py +4 -0
- haiku/rag/mcp.py +18 -6
- haiku/rag/qa/agent.py +4 -2
- haiku/rag/qa/prompts.py +2 -2
- haiku/rag/reranking/mxbai.py +1 -1
- haiku/rag/research/__init__.py +10 -27
- haiku/rag/research/common.py +53 -0
- haiku/rag/research/dependencies.py +3 -25
- haiku/rag/research/graph.py +29 -0
- haiku/rag/research/models.py +70 -0
- haiku/rag/research/nodes/evaluate.py +80 -0
- haiku/rag/research/nodes/plan.py +63 -0
- haiku/rag/research/nodes/search.py +93 -0
- haiku/rag/research/nodes/synthesize.py +51 -0
- haiku/rag/research/prompts.py +98 -113
- haiku/rag/research/state.py +25 -0
- haiku/rag/store/engine.py +14 -0
- haiku/rag/store/models/chunk.py +1 -0
- haiku/rag/store/models/document.py +1 -0
- haiku/rag/store/repositories/chunk.py +4 -0
- haiku/rag/store/repositories/document.py +3 -0
- haiku/rag/store/upgrades/__init__.py +2 -0
- haiku/rag/store/upgrades/v0_10_1.py +64 -0
- haiku/rag/utils.py +8 -5
- {haiku_rag-0.9.3.dist-info → haiku_rag-0.10.1.dist-info}/METADATA +37 -1
- haiku_rag-0.10.1.dist-info/RECORD +54 -0
- haiku/rag/research/base.py +0 -130
- haiku/rag/research/evaluation_agent.py +0 -85
- haiku/rag/research/orchestrator.py +0 -170
- haiku/rag/research/presearch_agent.py +0 -39
- haiku/rag/research/search_agent.py +0 -69
- haiku/rag/research/synthesis_agent.py +0 -60
- haiku_rag-0.9.3.dist-info/RECORD +0 -51
- {haiku_rag-0.9.3.dist-info → haiku_rag-0.10.1.dist-info}/WHEEL +0 -0
- {haiku_rag-0.9.3.dist-info → haiku_rag-0.10.1.dist-info}/entry_points.txt +0 -0
- {haiku_rag-0.9.3.dist-info → haiku_rag-0.10.1.dist-info}/licenses/LICENSE +0 -0
haiku/rag/research/base.py
DELETED
|
@@ -1,130 +0,0 @@
|
|
|
1
|
-
from abc import ABC, abstractmethod
|
|
2
|
-
from typing import TYPE_CHECKING, Any
|
|
3
|
-
|
|
4
|
-
from pydantic import BaseModel, Field
|
|
5
|
-
from pydantic_ai import Agent
|
|
6
|
-
from pydantic_ai.models.openai import OpenAIChatModel
|
|
7
|
-
from pydantic_ai.output import ToolOutput
|
|
8
|
-
from pydantic_ai.providers.ollama import OllamaProvider
|
|
9
|
-
from pydantic_ai.providers.openai import OpenAIProvider
|
|
10
|
-
from pydantic_ai.run import AgentRunResult
|
|
11
|
-
|
|
12
|
-
from haiku.rag.config import Config
|
|
13
|
-
|
|
14
|
-
if TYPE_CHECKING:
|
|
15
|
-
from haiku.rag.research.dependencies import ResearchDependencies
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
class BaseResearchAgent[T](ABC):
|
|
19
|
-
"""Base class for all research agents."""
|
|
20
|
-
|
|
21
|
-
def __init__(
|
|
22
|
-
self,
|
|
23
|
-
provider: str,
|
|
24
|
-
model: str,
|
|
25
|
-
output_type: type[T],
|
|
26
|
-
):
|
|
27
|
-
self.provider = provider
|
|
28
|
-
self.model = model
|
|
29
|
-
self.output_type = output_type
|
|
30
|
-
|
|
31
|
-
model_obj = self._get_model(provider, model)
|
|
32
|
-
|
|
33
|
-
# Import deps type lazily to avoid circular import during module load
|
|
34
|
-
from haiku.rag.research.dependencies import ResearchDependencies
|
|
35
|
-
|
|
36
|
-
# If the agent is expected to return plain text, pass `str` directly.
|
|
37
|
-
# Otherwise, wrap the model with ToolOutput for robust tool-handling retries.
|
|
38
|
-
agent_output_type: Any
|
|
39
|
-
if self.output_type is str: # plain text output
|
|
40
|
-
agent_output_type = str
|
|
41
|
-
else:
|
|
42
|
-
agent_output_type = ToolOutput(self.output_type, max_retries=3)
|
|
43
|
-
|
|
44
|
-
self._agent = Agent(
|
|
45
|
-
model=model_obj,
|
|
46
|
-
deps_type=ResearchDependencies,
|
|
47
|
-
output_type=agent_output_type,
|
|
48
|
-
instructions=self.get_system_prompt(),
|
|
49
|
-
retries=3,
|
|
50
|
-
)
|
|
51
|
-
|
|
52
|
-
# Register tools
|
|
53
|
-
self.register_tools()
|
|
54
|
-
|
|
55
|
-
def _get_model(self, provider: str, model: str):
|
|
56
|
-
"""Get the appropriate model object for the provider."""
|
|
57
|
-
if provider == "ollama":
|
|
58
|
-
return OpenAIChatModel(
|
|
59
|
-
model_name=model,
|
|
60
|
-
provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
|
|
61
|
-
)
|
|
62
|
-
elif provider == "vllm":
|
|
63
|
-
return OpenAIChatModel(
|
|
64
|
-
model_name=model,
|
|
65
|
-
provider=OpenAIProvider(
|
|
66
|
-
base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1",
|
|
67
|
-
api_key="none",
|
|
68
|
-
),
|
|
69
|
-
)
|
|
70
|
-
else:
|
|
71
|
-
# For all other providers, use the provider:model format
|
|
72
|
-
return f"{provider}:{model}"
|
|
73
|
-
|
|
74
|
-
@abstractmethod
|
|
75
|
-
def get_system_prompt(self) -> str:
|
|
76
|
-
"""Return the system prompt for this agent."""
|
|
77
|
-
pass
|
|
78
|
-
|
|
79
|
-
def register_tools(self) -> None:
|
|
80
|
-
"""Register agent-specific tools."""
|
|
81
|
-
pass
|
|
82
|
-
|
|
83
|
-
async def run(
|
|
84
|
-
self, prompt: str, deps: "ResearchDependencies", **kwargs
|
|
85
|
-
) -> AgentRunResult[T]:
|
|
86
|
-
"""Execute the agent."""
|
|
87
|
-
return await self._agent.run(prompt, deps=deps, **kwargs)
|
|
88
|
-
|
|
89
|
-
@property
|
|
90
|
-
def agent(self) -> Agent[Any, T]:
|
|
91
|
-
"""Access the underlying Pydantic AI agent."""
|
|
92
|
-
return self._agent
|
|
93
|
-
|
|
94
|
-
|
|
95
|
-
class SearchResult(BaseModel):
|
|
96
|
-
"""Standard search result format."""
|
|
97
|
-
|
|
98
|
-
content: str
|
|
99
|
-
score: float
|
|
100
|
-
document_uri: str
|
|
101
|
-
metadata: dict[str, Any] = Field(default_factory=dict)
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
class ResearchOutput(BaseModel):
|
|
105
|
-
"""Standard research output format."""
|
|
106
|
-
|
|
107
|
-
summary: str
|
|
108
|
-
detailed_findings: list[str]
|
|
109
|
-
sources: list[str]
|
|
110
|
-
confidence: float
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
class SearchAnswer(BaseModel):
|
|
114
|
-
"""Structured output for the SearchSpecialist agent."""
|
|
115
|
-
|
|
116
|
-
query: str = Field(description="The search query that was performed")
|
|
117
|
-
answer: str = Field(description="The answer generated based on the context")
|
|
118
|
-
context: list[str] = Field(
|
|
119
|
-
description=(
|
|
120
|
-
"Only the minimal set of relevant snippets (verbatim) that directly "
|
|
121
|
-
"support the answer"
|
|
122
|
-
)
|
|
123
|
-
)
|
|
124
|
-
sources: list[str] = Field(
|
|
125
|
-
description=(
|
|
126
|
-
"Document URIs corresponding to the snippets actually used in the"
|
|
127
|
-
" answer (one URI per snippet; omit if none)"
|
|
128
|
-
),
|
|
129
|
-
default_factory=list,
|
|
130
|
-
)
|
|
@@ -1,85 +0,0 @@
|
|
|
1
|
-
from pydantic import BaseModel, Field
|
|
2
|
-
from pydantic_ai.run import AgentRunResult
|
|
3
|
-
|
|
4
|
-
from haiku.rag.research.base import BaseResearchAgent
|
|
5
|
-
from haiku.rag.research.dependencies import (
|
|
6
|
-
ResearchDependencies,
|
|
7
|
-
_format_context_for_prompt,
|
|
8
|
-
)
|
|
9
|
-
from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT
|
|
10
|
-
|
|
11
|
-
|
|
12
|
-
class EvaluationResult(BaseModel):
|
|
13
|
-
"""Result of analysis and evaluation."""
|
|
14
|
-
|
|
15
|
-
key_insights: list[str] = Field(
|
|
16
|
-
description="Main insights extracted from the research so far"
|
|
17
|
-
)
|
|
18
|
-
new_questions: list[str] = Field(
|
|
19
|
-
description="New sub-questions to add to the research (max 3)",
|
|
20
|
-
max_length=3,
|
|
21
|
-
default=[],
|
|
22
|
-
)
|
|
23
|
-
confidence_score: float = Field(
|
|
24
|
-
description="Confidence level in the completeness of research (0-1)",
|
|
25
|
-
ge=0.0,
|
|
26
|
-
le=1.0,
|
|
27
|
-
)
|
|
28
|
-
is_sufficient: bool = Field(
|
|
29
|
-
description="Whether the research is sufficient to answer the original question"
|
|
30
|
-
)
|
|
31
|
-
reasoning: str = Field(
|
|
32
|
-
description="Explanation of why the research is or isn't complete"
|
|
33
|
-
)
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
class AnalysisEvaluationAgent(BaseResearchAgent[EvaluationResult]):
|
|
37
|
-
"""Agent that analyzes findings and evaluates research completeness."""
|
|
38
|
-
|
|
39
|
-
def __init__(self, provider: str, model: str) -> None:
|
|
40
|
-
super().__init__(provider, model, output_type=EvaluationResult)
|
|
41
|
-
|
|
42
|
-
async def run(
|
|
43
|
-
self, prompt: str, deps: ResearchDependencies, **kwargs
|
|
44
|
-
) -> AgentRunResult[EvaluationResult]:
|
|
45
|
-
console = deps.console
|
|
46
|
-
if console:
|
|
47
|
-
console.print(
|
|
48
|
-
"\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]"
|
|
49
|
-
)
|
|
50
|
-
|
|
51
|
-
# Format context for the evaluation agent
|
|
52
|
-
context_xml = _format_context_for_prompt(deps.context)
|
|
53
|
-
evaluation_prompt = f"""Analyze all gathered information and evaluate the completeness of research.
|
|
54
|
-
|
|
55
|
-
{context_xml}
|
|
56
|
-
|
|
57
|
-
Evaluate the research progress for the original question and identify any remaining gaps."""
|
|
58
|
-
|
|
59
|
-
result = await super().run(evaluation_prompt, deps, **kwargs)
|
|
60
|
-
output = result.output
|
|
61
|
-
|
|
62
|
-
# Store insights
|
|
63
|
-
for insight in output.key_insights:
|
|
64
|
-
deps.context.add_insight(insight)
|
|
65
|
-
|
|
66
|
-
# Add new questions to the sub-questions list
|
|
67
|
-
for new_q in output.new_questions:
|
|
68
|
-
if new_q not in deps.context.sub_questions:
|
|
69
|
-
deps.context.sub_questions.append(new_q)
|
|
70
|
-
|
|
71
|
-
if console:
|
|
72
|
-
if output.key_insights:
|
|
73
|
-
console.print(" [bold]Key insights:[/bold]")
|
|
74
|
-
for insight in output.key_insights:
|
|
75
|
-
console.print(f" • {insight}")
|
|
76
|
-
console.print(
|
|
77
|
-
f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]"
|
|
78
|
-
)
|
|
79
|
-
status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
|
|
80
|
-
console.print(f" Sufficient: {status}")
|
|
81
|
-
|
|
82
|
-
return result
|
|
83
|
-
|
|
84
|
-
def get_system_prompt(self) -> str:
|
|
85
|
-
return EVALUATION_AGENT_PROMPT
|
|
@@ -1,170 +0,0 @@
|
|
|
1
|
-
from typing import Any
|
|
2
|
-
|
|
3
|
-
from pydantic import BaseModel, Field
|
|
4
|
-
from pydantic_ai.run import AgentRunResult
|
|
5
|
-
from rich.console import Console
|
|
6
|
-
|
|
7
|
-
from haiku.rag.config import Config
|
|
8
|
-
from haiku.rag.research.base import BaseResearchAgent
|
|
9
|
-
from haiku.rag.research.dependencies import (
|
|
10
|
-
ResearchContext,
|
|
11
|
-
ResearchDependencies,
|
|
12
|
-
)
|
|
13
|
-
from haiku.rag.research.evaluation_agent import (
|
|
14
|
-
AnalysisEvaluationAgent,
|
|
15
|
-
EvaluationResult,
|
|
16
|
-
)
|
|
17
|
-
from haiku.rag.research.presearch_agent import PresearchSurveyAgent
|
|
18
|
-
from haiku.rag.research.prompts import ORCHESTRATOR_PROMPT
|
|
19
|
-
from haiku.rag.research.search_agent import SearchSpecialistAgent
|
|
20
|
-
from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
class ResearchPlan(BaseModel):
|
|
24
|
-
"""Research execution plan."""
|
|
25
|
-
|
|
26
|
-
main_question: str = Field(description="The main research question")
|
|
27
|
-
sub_questions: list[str] = Field(
|
|
28
|
-
description="Decomposed sub-questions to investigate (max 3)", max_length=3
|
|
29
|
-
)
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
|
|
33
|
-
"""Orchestrator agent that coordinates the research workflow."""
|
|
34
|
-
|
|
35
|
-
def __init__(
|
|
36
|
-
self,
|
|
37
|
-
provider: str | None = Config.RESEARCH_PROVIDER,
|
|
38
|
-
model: str | None = None,
|
|
39
|
-
):
|
|
40
|
-
# Use provided values or fall back to config defaults
|
|
41
|
-
provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER
|
|
42
|
-
model = model or Config.RESEARCH_MODEL or Config.QA_MODEL
|
|
43
|
-
|
|
44
|
-
super().__init__(provider, model, output_type=ResearchPlan)
|
|
45
|
-
|
|
46
|
-
self.search_agent: SearchSpecialistAgent = SearchSpecialistAgent(
|
|
47
|
-
provider, model
|
|
48
|
-
)
|
|
49
|
-
self.presearch_agent: PresearchSurveyAgent = PresearchSurveyAgent(
|
|
50
|
-
provider, model
|
|
51
|
-
)
|
|
52
|
-
self.evaluation_agent: AnalysisEvaluationAgent = AnalysisEvaluationAgent(
|
|
53
|
-
provider, model
|
|
54
|
-
)
|
|
55
|
-
self.synthesis_agent: SynthesisAgent = SynthesisAgent(provider, model)
|
|
56
|
-
|
|
57
|
-
def get_system_prompt(self) -> str:
|
|
58
|
-
return ORCHESTRATOR_PROMPT
|
|
59
|
-
|
|
60
|
-
def _should_stop_research(
|
|
61
|
-
self,
|
|
62
|
-
evaluation_result: AgentRunResult[EvaluationResult],
|
|
63
|
-
confidence_threshold: float,
|
|
64
|
-
) -> bool:
|
|
65
|
-
"""Determine if research should stop based on evaluation."""
|
|
66
|
-
|
|
67
|
-
result = evaluation_result.output
|
|
68
|
-
return result.is_sufficient and result.confidence_score >= confidence_threshold
|
|
69
|
-
|
|
70
|
-
async def conduct_research(
|
|
71
|
-
self,
|
|
72
|
-
question: str,
|
|
73
|
-
client: Any,
|
|
74
|
-
max_iterations: int = 3,
|
|
75
|
-
confidence_threshold: float = 0.8,
|
|
76
|
-
verbose: bool = False,
|
|
77
|
-
) -> ResearchReport:
|
|
78
|
-
"""Conduct comprehensive research on a question.
|
|
79
|
-
|
|
80
|
-
Args:
|
|
81
|
-
question: The research question to investigate
|
|
82
|
-
client: HaikuRAG client for document operations
|
|
83
|
-
max_iterations: Maximum number of search-analyze-clarify cycles
|
|
84
|
-
confidence_threshold: Minimum confidence level to stop research (0-1)
|
|
85
|
-
verbose: If True, print progress and intermediate results
|
|
86
|
-
|
|
87
|
-
Returns:
|
|
88
|
-
ResearchReport with comprehensive findings
|
|
89
|
-
"""
|
|
90
|
-
|
|
91
|
-
# Initialize context
|
|
92
|
-
context = ResearchContext(original_question=question)
|
|
93
|
-
deps = ResearchDependencies(client=client, context=context)
|
|
94
|
-
if verbose:
|
|
95
|
-
deps.console = Console()
|
|
96
|
-
|
|
97
|
-
console = deps.console
|
|
98
|
-
# Create initial research plan
|
|
99
|
-
if console:
|
|
100
|
-
console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
|
|
101
|
-
|
|
102
|
-
# Run a simple presearch survey to summarize KB context
|
|
103
|
-
presearch_result = await self.presearch_agent.run(question, deps=deps)
|
|
104
|
-
plan_prompt = (
|
|
105
|
-
"Create a research plan for the main question below.\n\n"
|
|
106
|
-
f"Main question: {question}\n\n"
|
|
107
|
-
"Use this brief presearch summary to inform the plan. Focus the 3 sub-questions "
|
|
108
|
-
"on the most important aspects not already obvious from the current KB context.\n\n"
|
|
109
|
-
f"{presearch_result.output}"
|
|
110
|
-
)
|
|
111
|
-
|
|
112
|
-
plan_result: AgentRunResult[ResearchPlan] = await self.run(
|
|
113
|
-
plan_prompt, deps=deps
|
|
114
|
-
)
|
|
115
|
-
context.sub_questions = plan_result.output.sub_questions
|
|
116
|
-
|
|
117
|
-
if console:
|
|
118
|
-
console.print("\n[bold green]✅ Research Plan Created:[/bold green]")
|
|
119
|
-
console.print(
|
|
120
|
-
f" [bold]Main Question:[/bold] {plan_result.output.main_question}"
|
|
121
|
-
)
|
|
122
|
-
console.print(" [bold]Sub-questions:[/bold]")
|
|
123
|
-
for i, sq in enumerate(plan_result.output.sub_questions, 1):
|
|
124
|
-
console.print(f" {i}. {sq}")
|
|
125
|
-
|
|
126
|
-
# Execute research iterations
|
|
127
|
-
for iteration in range(max_iterations):
|
|
128
|
-
if console:
|
|
129
|
-
console.rule(
|
|
130
|
-
f"[bold yellow]🔄 Iteration {iteration + 1}/{max_iterations}[/bold yellow]"
|
|
131
|
-
)
|
|
132
|
-
|
|
133
|
-
# Check if we have questions to search
|
|
134
|
-
if not context.sub_questions:
|
|
135
|
-
if console:
|
|
136
|
-
console.print(
|
|
137
|
-
"[yellow]No more questions to explore. Concluding research.[/yellow]"
|
|
138
|
-
)
|
|
139
|
-
break
|
|
140
|
-
|
|
141
|
-
# Use current sub-questions for this iteration
|
|
142
|
-
questions_to_search = context.sub_questions[:]
|
|
143
|
-
|
|
144
|
-
# Search phase - answer all questions in this iteration
|
|
145
|
-
if console:
|
|
146
|
-
console.print(
|
|
147
|
-
f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]"
|
|
148
|
-
)
|
|
149
|
-
|
|
150
|
-
for search_question in questions_to_search:
|
|
151
|
-
await self.search_agent.run(search_question, deps=deps)
|
|
152
|
-
|
|
153
|
-
# Analysis and Evaluation phase
|
|
154
|
-
|
|
155
|
-
evaluation_result = await self.evaluation_agent.run("", deps=deps)
|
|
156
|
-
|
|
157
|
-
# Check if research is sufficient
|
|
158
|
-
if self._should_stop_research(evaluation_result, confidence_threshold):
|
|
159
|
-
if console:
|
|
160
|
-
console.print(
|
|
161
|
-
f"\n[bold green]✅ Stopping research:[/bold green] {evaluation_result.output.reasoning}"
|
|
162
|
-
)
|
|
163
|
-
break
|
|
164
|
-
|
|
165
|
-
# Generate final report
|
|
166
|
-
report_result: AgentRunResult[ResearchReport] = await self.synthesis_agent.run(
|
|
167
|
-
"", deps=deps
|
|
168
|
-
)
|
|
169
|
-
|
|
170
|
-
return report_result.output
|
|
@@ -1,39 +0,0 @@
|
|
|
1
|
-
from pydantic_ai import RunContext
|
|
2
|
-
from pydantic_ai.run import AgentRunResult
|
|
3
|
-
|
|
4
|
-
from haiku.rag.research.base import BaseResearchAgent
|
|
5
|
-
from haiku.rag.research.dependencies import ResearchDependencies
|
|
6
|
-
from haiku.rag.research.prompts import PRESEARCH_AGENT_PROMPT
|
|
7
|
-
|
|
8
|
-
|
|
9
|
-
class PresearchSurveyAgent(BaseResearchAgent[str]):
|
|
10
|
-
"""Presearch agent that gathers verbatim context and summarizes it."""
|
|
11
|
-
|
|
12
|
-
def __init__(self, provider: str, model: str) -> None:
|
|
13
|
-
super().__init__(provider, model, str)
|
|
14
|
-
|
|
15
|
-
async def run(
|
|
16
|
-
self, prompt: str, deps: ResearchDependencies, **kwargs
|
|
17
|
-
) -> AgentRunResult[str]:
|
|
18
|
-
console = deps.console
|
|
19
|
-
if console:
|
|
20
|
-
console.print(
|
|
21
|
-
"\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]"
|
|
22
|
-
)
|
|
23
|
-
|
|
24
|
-
return await super().run(prompt, deps, **kwargs)
|
|
25
|
-
|
|
26
|
-
def get_system_prompt(self) -> str:
|
|
27
|
-
return PRESEARCH_AGENT_PROMPT
|
|
28
|
-
|
|
29
|
-
def register_tools(self) -> None:
|
|
30
|
-
@self.agent.tool
|
|
31
|
-
async def gather_context(
|
|
32
|
-
ctx: RunContext[ResearchDependencies],
|
|
33
|
-
query: str,
|
|
34
|
-
limit: int = 6,
|
|
35
|
-
) -> str:
|
|
36
|
-
"""Return verbatim concatenation of relevant chunk texts."""
|
|
37
|
-
results = await ctx.deps.client.search(query, limit=limit)
|
|
38
|
-
expanded = await ctx.deps.client.expand_context(results)
|
|
39
|
-
return "\n\n".join(chunk.content for chunk, _ in expanded)
|
|
@@ -1,69 +0,0 @@
|
|
|
1
|
-
from pydantic_ai import RunContext
|
|
2
|
-
from pydantic_ai.format_prompt import format_as_xml
|
|
3
|
-
from pydantic_ai.run import AgentRunResult
|
|
4
|
-
|
|
5
|
-
from haiku.rag.research.base import BaseResearchAgent, SearchAnswer
|
|
6
|
-
from haiku.rag.research.dependencies import ResearchDependencies
|
|
7
|
-
from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT
|
|
8
|
-
|
|
9
|
-
|
|
10
|
-
class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
|
|
11
|
-
"""Agent specialized in answering questions using RAG search."""
|
|
12
|
-
|
|
13
|
-
def __init__(self, provider: str, model: str) -> None:
|
|
14
|
-
super().__init__(provider, model, output_type=SearchAnswer)
|
|
15
|
-
|
|
16
|
-
async def run(
|
|
17
|
-
self, prompt: str, deps: ResearchDependencies, **kwargs
|
|
18
|
-
) -> AgentRunResult[SearchAnswer]:
|
|
19
|
-
"""Execute the agent and persist the QA pair in shared context.
|
|
20
|
-
|
|
21
|
-
Pydantic AI enforces `SearchAnswer` as the output model; we just store
|
|
22
|
-
the QA response with the last search results as sources.
|
|
23
|
-
"""
|
|
24
|
-
console = deps.console
|
|
25
|
-
if console:
|
|
26
|
-
console.print(f"\t{prompt}")
|
|
27
|
-
|
|
28
|
-
result = await super().run(prompt, deps, **kwargs)
|
|
29
|
-
deps.context.add_qa_response(result.output)
|
|
30
|
-
deps.context.sub_questions.remove(prompt)
|
|
31
|
-
if console:
|
|
32
|
-
answer = result.output.answer
|
|
33
|
-
answer_preview = answer[:150] + "…" if len(answer) > 150 else answer
|
|
34
|
-
console.log(f"\n [green]✓[/green] {answer_preview}")
|
|
35
|
-
|
|
36
|
-
return result
|
|
37
|
-
|
|
38
|
-
def get_system_prompt(self) -> str:
|
|
39
|
-
return SEARCH_AGENT_PROMPT
|
|
40
|
-
|
|
41
|
-
def register_tools(self) -> None:
|
|
42
|
-
"""Register search-specific tools."""
|
|
43
|
-
|
|
44
|
-
@self.agent.tool
|
|
45
|
-
async def search_and_answer(
|
|
46
|
-
ctx: RunContext[ResearchDependencies],
|
|
47
|
-
query: str,
|
|
48
|
-
limit: int = 5,
|
|
49
|
-
) -> str:
|
|
50
|
-
"""Search the KB and return a concise context pack."""
|
|
51
|
-
search_results = await ctx.deps.client.search(query, limit=limit)
|
|
52
|
-
expanded = await ctx.deps.client.expand_context(search_results)
|
|
53
|
-
|
|
54
|
-
snippet_entries = [
|
|
55
|
-
{
|
|
56
|
-
"text": chunk.content,
|
|
57
|
-
"score": score,
|
|
58
|
-
"document_uri": (chunk.document_uri or ""),
|
|
59
|
-
}
|
|
60
|
-
for chunk, score in expanded
|
|
61
|
-
]
|
|
62
|
-
|
|
63
|
-
# Return an XML-formatted payload with the question and snippets.
|
|
64
|
-
if snippet_entries:
|
|
65
|
-
return format_as_xml(snippet_entries, root_tag="snippets")
|
|
66
|
-
else:
|
|
67
|
-
return (
|
|
68
|
-
f"No relevant information found in the knowledge base for: {query}"
|
|
69
|
-
)
|
|
@@ -1,60 +0,0 @@
|
|
|
1
|
-
from pydantic import BaseModel, Field
|
|
2
|
-
from pydantic_ai.run import AgentRunResult
|
|
3
|
-
|
|
4
|
-
from haiku.rag.research.base import BaseResearchAgent
|
|
5
|
-
from haiku.rag.research.dependencies import (
|
|
6
|
-
ResearchDependencies,
|
|
7
|
-
_format_context_for_prompt,
|
|
8
|
-
)
|
|
9
|
-
from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT
|
|
10
|
-
|
|
11
|
-
|
|
12
|
-
class ResearchReport(BaseModel):
|
|
13
|
-
"""Final research report structure."""
|
|
14
|
-
|
|
15
|
-
title: str = Field(description="Concise title for the research")
|
|
16
|
-
executive_summary: str = Field(description="Brief overview of key findings")
|
|
17
|
-
main_findings: list[str] = Field(
|
|
18
|
-
description="Primary research findings with supporting evidence"
|
|
19
|
-
)
|
|
20
|
-
conclusions: list[str] = Field(description="Evidence-based conclusions")
|
|
21
|
-
limitations: list[str] = Field(
|
|
22
|
-
description="Limitations of the current research", default=[]
|
|
23
|
-
)
|
|
24
|
-
recommendations: list[str] = Field(
|
|
25
|
-
description="Actionable recommendations based on findings", default=[]
|
|
26
|
-
)
|
|
27
|
-
sources_summary: str = Field(
|
|
28
|
-
description="Summary of sources used and their reliability"
|
|
29
|
-
)
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
class SynthesisAgent(BaseResearchAgent[ResearchReport]):
|
|
33
|
-
"""Agent specialized in synthesizing research into comprehensive reports."""
|
|
34
|
-
|
|
35
|
-
def __init__(self, provider: str, model: str) -> None:
|
|
36
|
-
super().__init__(provider, model, output_type=ResearchReport)
|
|
37
|
-
|
|
38
|
-
async def run(
|
|
39
|
-
self, prompt: str, deps: ResearchDependencies, **kwargs
|
|
40
|
-
) -> AgentRunResult[ResearchReport]:
|
|
41
|
-
console = deps.console
|
|
42
|
-
if console:
|
|
43
|
-
console.print(
|
|
44
|
-
"\n[bold cyan]📝 Generating final research report...[/bold cyan]"
|
|
45
|
-
)
|
|
46
|
-
|
|
47
|
-
context_xml = _format_context_for_prompt(deps.context)
|
|
48
|
-
synthesis_prompt = f"""Generate a comprehensive research report based on all gathered information.
|
|
49
|
-
|
|
50
|
-
{context_xml}
|
|
51
|
-
|
|
52
|
-
Create a detailed report that synthesizes all findings into a coherent response."""
|
|
53
|
-
result = await super().run(synthesis_prompt, deps, **kwargs)
|
|
54
|
-
if console:
|
|
55
|
-
console.print("[bold green]✅ Research complete![/bold green]")
|
|
56
|
-
|
|
57
|
-
return result
|
|
58
|
-
|
|
59
|
-
def get_system_prompt(self) -> str:
|
|
60
|
-
return SYNTHESIS_AGENT_PROMPT
|
haiku_rag-0.9.3.dist-info/RECORD
DELETED
|
@@ -1,51 +0,0 @@
|
|
|
1
|
-
haiku/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
2
|
-
haiku/rag/app.py,sha256=nkud-OHic3HIgEEiNOKVvhmW98DPpDe6HokBSz-xV7w,11420
|
|
3
|
-
haiku/rag/chunker.py,sha256=PVe6ysv8UlacUd4Zb3_8RFWIaWDXnzBAy2VDJ4TaUsE,1555
|
|
4
|
-
haiku/rag/cli.py,sha256=3nlzrT5FPCyfnu51KHchLG4Cj2eVv9YsuGHMShBnVb0,9845
|
|
5
|
-
haiku/rag/client.py,sha256=QgJQu7g7JjAzWN6R10NeDqpFf89Dml_LiWce4QRHLHc,21177
|
|
6
|
-
haiku/rag/config.py,sha256=SPEIv2IElZmZh4Wsp8gk7ViRW5ZzD-UGmIqRAXscDdI,2134
|
|
7
|
-
haiku/rag/logging.py,sha256=dm65AwADpcQsH5OAPtRA-4hsw0w5DK-sGOvzYkj6jzw,1720
|
|
8
|
-
haiku/rag/mcp.py,sha256=bR9Y-Nz-hvjiql20Y0KE0hwNGwyjmPGX8K9d-qmXptY,4683
|
|
9
|
-
haiku/rag/migration.py,sha256=M--KnSF3lxgKjxmokb4vuzGH-pV8eg0C_8e7jvPqW8Y,11058
|
|
10
|
-
haiku/rag/monitor.py,sha256=r386nkhdlsU8UECwIuVwnrSlgMk3vNIuUZGNIzkZuec,2770
|
|
11
|
-
haiku/rag/reader.py,sha256=qkPTMJuQ_o4sK-8zpDl9WFYe_MJ7aL_gUw6rczIpW-g,3274
|
|
12
|
-
haiku/rag/utils.py,sha256=aiuPu_rrfpyIvJJq0o5boUIIvCdNzdpKwAIPYYn3iG8,4965
|
|
13
|
-
haiku/rag/embeddings/__init__.py,sha256=44IfDITGIFTflGT6UEmiYOwpWFVbYv5smLY59D0YeCs,1419
|
|
14
|
-
haiku/rag/embeddings/base.py,sha256=BnSviKrlzjv3L0sZJs_T-pxfawd-bcTak-rsX-D2f3A,497
|
|
15
|
-
haiku/rag/embeddings/ollama.py,sha256=LuLlHH6RGoO9_gFCIlbmesuXOj017gTw6z-p8Ez0CfE,595
|
|
16
|
-
haiku/rag/embeddings/openai.py,sha256=fIFCk-jpUtaW0xsnrQnJ824O0UCjaGG2sgvBzREhilc,503
|
|
17
|
-
haiku/rag/embeddings/vllm.py,sha256=vhaUnCn6VMkfSluLhWKtSV-sekFaPsp4pKo2N7-SBCY,626
|
|
18
|
-
haiku/rag/embeddings/voyageai.py,sha256=UW-MW4tJKnPB6Fs2P7A3yt-ZeRm46H9npckchSriPX8,661
|
|
19
|
-
haiku/rag/qa/__init__.py,sha256=Sl7Kzrg9CuBOcMF01wc1NtQhUNWjJI0MhIHfCWrb8V4,434
|
|
20
|
-
haiku/rag/qa/agent.py,sha256=f4Keh-ESgctNbTg96QL95HYjINVLOcxa8t8crx92MMk,3081
|
|
21
|
-
haiku/rag/qa/prompts.py,sha256=LhRfDtO8Pb06lpr4PpwEaKUYItZ5OiIkeqcCogcssHY,3347
|
|
22
|
-
haiku/rag/reranking/__init__.py,sha256=IRXHs4qPu6VbGJQpzSwhgtVWWumURH_vEoVFE-extlo,894
|
|
23
|
-
haiku/rag/reranking/base.py,sha256=LM9yUSSJ414UgBZhFTgxGprlRqzfTe4I1vgjricz2JY,405
|
|
24
|
-
haiku/rag/reranking/cohere.py,sha256=1iTdiaa8vvb6oHVB2qpWzUOVkyfUcimVSZp6Qr4aq4c,1049
|
|
25
|
-
haiku/rag/reranking/mxbai.py,sha256=46sVTsTIkzIX9THgM3u8HaEmgY7evvEyB-N54JTHvK8,867
|
|
26
|
-
haiku/rag/reranking/vllm.py,sha256=xVGH9ss-ISWdJ5SKUUHUbTqBo7PIEmA_SQv0ScdJ6XA,1479
|
|
27
|
-
haiku/rag/research/__init__.py,sha256=qLF41YayAxW_VeHhuTceVuz9hw1FNbuRV9VMhonUMW0,1078
|
|
28
|
-
haiku/rag/research/base.py,sha256=X5n6myUG_Oz4i8WGfyKZ39YzK13rOkyvwGKwSBfL50k,4043
|
|
29
|
-
haiku/rag/research/dependencies.py,sha256=N7mnFwa_uyWYH0NtbEHp5JJvNGN64Q8HHfY41E8Irx0,2362
|
|
30
|
-
haiku/rag/research/evaluation_agent.py,sha256=VMegemd9Vln3jfZbeHzMfb7rUPFNzNxi5Y_l1zrddl8,2994
|
|
31
|
-
haiku/rag/research/orchestrator.py,sha256=nvSRdIs77kSb1CZaQUYZM_Zl5xLP8K6noVgnixpeLJI,6329
|
|
32
|
-
haiku/rag/research/presearch_agent.py,sha256=MpakZ9HSynv73EnWakwUuytfKpiN_8lEqZlVc3zZjGU,1427
|
|
33
|
-
haiku/rag/research/prompts.py,sha256=pVRB7_b_p3JaLF1bC3ANTbSFY78ypSjDhoq6peoU6jo,5685
|
|
34
|
-
haiku/rag/research/search_agent.py,sha256=xn2MlEyL9te_dtZqTzW81lGw7fYmyUzn26mvzX52hNA,2599
|
|
35
|
-
haiku/rag/research/synthesis_agent.py,sha256=FQCt8wbaaKOwgIOQazTNAmohBMZRUDoVzHkByYhbGg8,2182
|
|
36
|
-
haiku/rag/store/__init__.py,sha256=hq0W0DAC7ysqhWSP2M2uHX8cbG6kbr-sWHxhq6qQcY0,103
|
|
37
|
-
haiku/rag/store/engine.py,sha256=-3MZJYft2XTWaLuyKha8DKhWQeU5E5CBeskXXF5fXso,9555
|
|
38
|
-
haiku/rag/store/models/__init__.py,sha256=s0E72zneGlowvZrFWaNxHYjOAUjgWdLxzdYsnvNRVlY,88
|
|
39
|
-
haiku/rag/store/models/chunk.py,sha256=Ww_hj3DMwJLNM33l1GvIP84yzDFc6cxfiWcotUfWSYg,383
|
|
40
|
-
haiku/rag/store/models/document.py,sha256=zSSpt6pyrMJAIXGQvIcqojcqUzwZnhp3WxVokaWxNRc,396
|
|
41
|
-
haiku/rag/store/repositories/__init__.py,sha256=Olv5dLfBQINRV3HrsfUpjzkZ7Qm7goEYyMNykgo_DaY,291
|
|
42
|
-
haiku/rag/store/repositories/chunk.py,sha256=O2SEhQy3ZptWjwwpxS-L8KNq2tEqEBqheHfLw-M_FqA,15012
|
|
43
|
-
haiku/rag/store/repositories/document.py,sha256=m11SamQoGYs5ODfmarJGU1yIcqtgmnba-5bGOPQuYrI,7773
|
|
44
|
-
haiku/rag/store/repositories/settings.py,sha256=7XMBMavU8zRgdBoQzQg0Obfa7UKjuVnBugidTC6sEW0,5548
|
|
45
|
-
haiku/rag/store/upgrades/__init__.py,sha256=gDOxiq3wdZPr3JoenjNYxx0cpgZJhbaFKNX2fzXRq1Q,1852
|
|
46
|
-
haiku/rag/store/upgrades/v0_9_3.py,sha256=NrjNilQSgDtFWRbL3ZUtzQzJ8tf9u0dDRJtnDFwwbdw,3322
|
|
47
|
-
haiku_rag-0.9.3.dist-info/METADATA,sha256=iCae4YtZ0meIQTZLUNree_-74F3irXvArPvdSxVz8ZM,4681
|
|
48
|
-
haiku_rag-0.9.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
|
49
|
-
haiku_rag-0.9.3.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
|
|
50
|
-
haiku_rag-0.9.3.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
|
|
51
|
-
haiku_rag-0.9.3.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|
|
File without changes
|