arsam
b3b70f9ce4
Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.
118 lines
3.9 KiB
Python
118 lines
3.9 KiB
Python
import logging
|
|
|
|
import httpx
|
|
|
|
from axon.config import settings
|
|
from axon.models import SearchResult
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_SYSTEM_PROMPT = (
|
|
"You are a precise research assistant. Answer the user's query based only on the "
|
|
"provided search results. Cite sources using [1], [2] notation. Be concise and factual."
|
|
)
|
|
|
|
|
|
def _headers() -> dict:
|
|
h = {"Content-Type": "application/json"}
|
|
if settings.litellm_api_key:
|
|
h["Authorization"] = f"Bearer {settings.litellm_api_key}"
|
|
return h
|
|
|
|
|
|
def _context_block(results: list[SearchResult], limit: int = 20) -> str:
|
|
parts = []
|
|
for i, r in enumerate(results[:limit], 1):
|
|
parts.append(f"[{i}] {r.title}\nURL: {r.url}\n{r.content}")
|
|
return "\n\n".join(parts)
|
|
|
|
|
|
async def _chat(messages: list[dict], max_tokens: int = 1024, timeout: float = 30.0) -> str:
|
|
payload = {
|
|
"model": settings.synthesis_model,
|
|
"messages": messages,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
async with httpx.AsyncClient() as client:
|
|
response = await client.post(
|
|
f"{settings.litellm_base_url}/chat/completions",
|
|
json=payload,
|
|
headers=_headers(),
|
|
timeout=timeout,
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()["choices"][0]["message"]["content"]
|
|
|
|
|
|
async def synthesize_answer(query: str, results: list[SearchResult]) -> str:
|
|
if not results:
|
|
return "No search results available to synthesize an answer."
|
|
try:
|
|
return await _chat(
|
|
[
|
|
{"role": "system", "content": _SYSTEM_PROMPT},
|
|
{"role": "user", "content": f"Query: {query}\n\nSearch Results:\n{_context_block(results, 5)}"},
|
|
]
|
|
)
|
|
except Exception as exc:
|
|
logger.error("Answer synthesis failed: %s", exc)
|
|
return f"Answer synthesis failed: {exc}"
|
|
|
|
|
|
async def generate_sub_queries(query: str, count: int = 3) -> list[str]:
|
|
try:
|
|
text = await _chat(
|
|
[
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
f"Generate exactly {count} distinct search sub-queries to research the given topic. "
|
|
"Each sub-query explores a different angle. "
|
|
"Return ONLY the queries, one per line, no numbering or explanation."
|
|
),
|
|
},
|
|
{"role": "user", "content": query},
|
|
],
|
|
max_tokens=256,
|
|
)
|
|
queries = [q.strip() for q in text.strip().splitlines() if q.strip()]
|
|
return queries[:count]
|
|
except Exception as exc:
|
|
logger.error("Sub-query generation failed: %s", exc)
|
|
return [query, f"{query} overview", f"{query} recent developments"][:count]
|
|
|
|
|
|
async def synthesize_report(
|
|
query: str,
|
|
results: list[SearchResult],
|
|
output_format: str = "markdown",
|
|
) -> str:
|
|
if not results:
|
|
return "No search results available for report generation."
|
|
|
|
format_instruction = (
|
|
"Format as valid JSON with keys: summary (string), key_findings (list of strings), conclusions (string)."
|
|
if output_format == "json"
|
|
else "Format as clear markdown with sections: ## Summary, ## Key Findings, ## Conclusion."
|
|
)
|
|
|
|
try:
|
|
return await _chat(
|
|
[
|
|
{"role": "system", "content": _SYSTEM_PROMPT},
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
f"Topic: {query}\n\nSearch Results:\n{_context_block(results)}\n\n"
|
|
f"Write a comprehensive research report. {format_instruction} "
|
|
"Cite sources using [N] notation."
|
|
),
|
|
},
|
|
],
|
|
max_tokens=2048,
|
|
timeout=60.0,
|
|
)
|
|
except Exception as exc:
|
|
logger.error("Report synthesis failed: %s", exc)
|
|
return f"Report synthesis failed: {exc}"
|