diff --git a/axon/services/synthesizer.py b/axon/services/synthesizer.py new file mode 100644 index 0000000..2ede576 --- /dev/null +++ b/axon/services/synthesizer.py @@ -0,0 +1,117 @@ +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}"