diff --git a/ai_search_agent/agent.py b/ai_search_agent/agent.py index 5eb3b62..a3d7b18 100644 --- a/ai_search_agent/agent.py +++ b/ai_search_agent/agent.py @@ -1,23 +1,346 @@ from __future__ import annotations -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator +from datetime import datetime, timezone +from .cache import SearchCache from .config import Settings from .jina import JinaClient +from .knowledge_graph import KnowledgeGraphExtractor from .llm_client import LLMClient from .models import ( AnswerPayload, Citation, + KnowledgeTriple, PageContent, SearchRequest, SearchResponse, SearchResult, SearchRoundTrace, ) +from .parsers import get_domain_trust +from .planner import QueryPlanner from .prompts import SYSTEM_PROMPT +from .reranker import JinaReranker from .search_client import SearchClient +class AISearchAgent: + def __init__(self, settings: Settings): + self.settings = settings + self.jina = JinaClient(settings) + self.llm = LLMClient(settings) + self.search = SearchClient(settings) + self.cache = SearchCache(settings) + self.reranker = JinaReranker(settings) + self.planner = QueryPlanner(settings) + self.kg = KnowledgeGraphExtractor(settings) + + # ------------------------------------------------------------------ # + # Public interface # + # ------------------------------------------------------------------ # + + async def run(self, request: SearchRequest) -> SearchResponse: + result: SearchResponse | None = None + async for event in self.run_stream(request): + if event["type"] == "result": + result = SearchResponse.model_validate(event["data"]) + if result is None: + raise RuntimeError("Search run completed without result payload") + return result + + async def run_stream( + self, request: SearchRequest + ) -> AsyncGenerator[dict, None]: + top_k = request.top_k_pages or self.settings.top_k_pages + max_chars = request.max_page_chars or self.settings.max_page_chars + + # ── 0. Cache hit ──────────────────────────────────────────────── + if self.settings.enable_cache: + cached = await self.cache.get_search(request.query, top_k, max_chars) + if cached: + cached["cache_hit"] = True + yield {"type": "cache_hit", "data": {"query": request.query}} + yield {"type": "result", "data": cached} + return + + # ── 1. History memory ──────────────────────────────────────────── + history_context = "" + if self.settings.enable_memory: + history = await self.cache.get_history(5) + recent = [h.get("query", "") for h in history if h.get("query")] + if recent: + history_context = "User's recent searches: " + "; ".join(recent[:5]) + + # ── 2. Query decomposition ─────────────────────────────────────── + sub_queries = await self.planner.decompose(request.query) + yield {"type": "decomposed", "data": {"sub_queries": sub_queries}} + + # ── 3. Multilingual expansion ──────────────────────────────────── + expanded: list[str] = [] + seen_q: set[str] = set() + for sq in sub_queries: + for q in await self.planner.expand_multilingual(sq): + if q not in seen_q: + seen_q.add(q) + expanded.append(q) + if len(expanded) > len(sub_queries): + yield {"type": "expanded", "data": {"queries": expanded}} + + # ── 4. Multi-round search loop ─────────────────────────────────── + all_results: list[SearchResult] = [] + all_pages: list[PageContent] = [] + all_usable: list[PageContent] = [] + round_traces: list[SearchRoundTrace] = [] + latest_digest = "" + seen_urls: set[str] = set() + current_queries = expanded + max_rounds = self.settings.max_search_rounds + round_index = 0 + + for round_index in range(1, max_rounds + 1): + yield { + "type": "round_started", + "data": {"round_index": round_index, "queries": current_queries}, + } + + round_results: list[SearchResult] = [] + for q in current_queries: + try: + raw_text, results = await self.search.search(q) + latest_digest = raw_text + for r in results: + if r.url not in seen_urls: + seen_urls.add(r.url) + round_results.append(r) + except Exception as exc: + yield {"type": "search_error", "data": {"query": q, "error": str(exc)}} + + pages = await self.jina.read_pages(round_results, max_page_chars=max_chars) + usable = [p for p in pages if p.fetched and p.usable] + + if self.settings.enable_reranker and len(usable) > 1: + try: + usable = await self.reranker.rerank(request.query, usable) + except Exception: + pass + + reflection = self._reflect( + round_results, pages, usable, round_index, max_rounds, + len(all_usable) + len(usable), + ) + trace = SearchRoundTrace( + round_index=round_index, + query=current_queries[0] if current_queries else request.query, + queries=current_queries, + result_count=len(round_results), + fetched_page_count=len([p for p in pages if p.fetched]), + usable_page_count=len(usable), + reflection=reflection, + ) + round_traces.append(trace) + all_results.extend(round_results) + all_pages.extend(pages) + all_usable.extend(usable) + + yield {"type": "round_finished", "data": trace.model_dump()} + + if len(all_usable) >= top_k or round_index >= max_rounds: + break + + next_q = await self.planner.next_query( + request.query, round_index, latest_digest, len(all_usable) + ) + if not next_q or next_q in seen_q: + break + seen_q.add(next_q) + current_queries = [next_q] + + # ── 5. No usable content ───────────────────────────────────────── + if not all_usable: + response = SearchResponse( + query=request.query, + answer=AnswerPayload( + summary="未读取到可用正文内容,已拒绝调用模型总结。", + caveats=[ + "所有候选页面都未能提取到有效正文,可能是 403、451、反爬或页面过短。", + "请更换查询词,或增加可抓取来源。", + ], + ), + search_results=all_results, + pages=all_pages, + search_rounds=round_traces, + content_ready=False, + llm_called=False, + error="NO_USABLE_CONTENT", + ) + yield {"type": "result", "data": response.model_dump()} + return + + # ── 6. Prepare model inputs ────────────────────────────────────── + pages_for_model = all_usable[:top_k] + results_for_model = [ + r for r in all_results if any(p.url == r.url for p in pages_for_model) + ] + allowed_urls = {p.url for p in pages_for_model if p.fetched} + user_prompt = self._build_user_prompt( + request.query, latest_digest, results_for_model, pages_for_model, history_context + ) + + # ── 7. Token-level LLM streaming ───────────────────────────────── + yield { + "type": "llm_started", + "data": {"page_count": len(pages_for_model), "result_count": len(results_for_model)}, + } + + answer: AnswerPayload | None = None + async for chunk, final_answer in self.llm.answer_stream(SYSTEM_PROMPT, user_prompt, allowed_urls): + if final_answer is not None: + answer = final_answer + elif chunk: + yield {"type": "llm_token", "data": {"text": chunk}} + + if answer is None: + answer = AnswerPayload(summary="", caveats=["LLM 未返回有效响应。"]) + + # ── 8. Confidence-based auto re-search ─────────────────────────── + if answer.confidence < self.settings.confidence_threshold: + extra_q = await self.planner.next_query( + request.query, round_index + 1, latest_digest, len(all_usable) + ) + if extra_q and extra_q not in seen_q: + yield { + "type": "confidence_retry", + "data": {"confidence": answer.confidence, "extra_query": extra_q}, + } + try: + _, extra_results = await self.search.search(extra_q) + new_results = [r for r in extra_results if r.url not in seen_urls] + if new_results: + extra_pages = await self.jina.read_pages( + new_results[:3], max_page_chars=max_chars + ) + extra_usable = [p for p in extra_pages if p.fetched and p.usable] + if extra_usable: + pages_for_model = (pages_for_model + extra_usable)[: top_k + 2] + allowed_urls.update(p.url for p in extra_usable) + user_prompt2 = self._build_user_prompt( + request.query, latest_digest, + results_for_model, pages_for_model, history_context, + ) + answer = await self.llm.answer(SYSTEM_PROMPT, user_prompt2, allowed_urls) + except Exception: + pass + + # ── 9. Backfill citations + domain trust scoring ───────────────── + if not answer.citations: + answer = self._backfill_citations(answer, results_for_model) + answer = self._score_citations(answer) + + # ── 10. KG extraction ──────────────────────────────────────────── + knowledge_triples: list[KnowledgeTriple] = [] + if self.settings.enable_kg: + combined_text = "\n".join(p.content for p in pages_for_model if p.content) + try: + knowledge_triples = await self.kg.extract(combined_text, request.query) + except Exception: + pass + + # ── 11. Assemble response ──────────────────────────────────────── + response = SearchResponse( + query=request.query, + answer=answer, + search_results=results_for_model, + pages=all_pages, + search_rounds=round_traces, + knowledge_triples=knowledge_triples, + content_ready=True, + llm_called=True, + ) + + # ── 12. Cache + history ────────────────────────────────────────── + if self.settings.enable_cache: + await self.cache.set_search(request.query, top_k, max_chars, response.model_dump()) + if self.settings.enable_memory: + await self.cache.push_history( + { + "query": request.query, + "summary": answer.summary[:100], + "ts": datetime.now(timezone.utc).isoformat(), + } + ) + + yield {"type": "result", "data": response.model_dump()} + + # ------------------------------------------------------------------ # + # Helpers # + # ------------------------------------------------------------------ # + + def _reflect( + self, + results: list[SearchResult], + pages: list[PageContent], + usable: list[PageContent], + round_index: int, + max_rounds: int, + total_usable: int, + ) -> str: + if not results: + return "未解析到搜索结果,模型将规划新查询词。" + if not pages: + return "未读取到任何页面内容,下一轮扩大来源。" + if total_usable >= self.settings.top_k_pages: + return "可用证据已达阈值,进入模型总结。" + if round_index >= max_rounds: + return "达到最大搜索轮次,使用当前证据总结。" + blocked = sum(1 for p in pages if p.fetched and not p.usable) + if blocked: + return f"本轮 {blocked} 个页面疑似反爬或正文不足,下一轮将换方向。" + return "证据仍不足,模型将规划下一轮查询。" + + def _build_user_prompt( + self, + query: str, + raw_search_text: str, + search_results: list[SearchResult], + pages: list[PageContent], + history_context: str = "", + ) -> str: + result_lines: list[str] = [] + for r in search_results: + result_lines.append(f"[Result {r.rank}] {r.title or 'Untitled'}") + result_lines.append(f"URL: {r.url}") + if r.snippet: + result_lines.append(f"Snippet: {r.snippet}") + + page_lines: list[str] = [] + for i, page in enumerate(pages, 1): + page_lines.append(f"[Page {i}] {page.title or page.url}") + page_lines.append(f"URL: {page.url}") + page_lines.append(page.content if page.fetched else f"Fetch error: {page.error}") + + ctx = f"\nUser context: {history_context}\n" if history_context else "" + return ( + f"User query:\n{query}{ctx}\n\n" + f"Search digest:\n{raw_search_text[:2000]}\n\n" + f"Selected search results:\n{chr(10).join(result_lines)}\n\n" + f"Fetched page excerpts:\n{chr(10).join(page_lines)}" + ) + + def _backfill_citations( + self, answer: AnswerPayload, results: list[SearchResult] + ) -> AnswerPayload: + fallback = [Citation(title=r.title or r.url, url=r.url) for r in results[:3]] + return answer.model_copy(update={"citations": fallback}) + + def _score_citations(self, answer: AnswerPayload) -> AnswerPayload: + scored = [ + c.model_copy(update={"trust_score": get_domain_trust(c.url)}) + for c in answer.citations + ] + return answer.model_copy(update={"citations": scored}) + + class AISearchAgent: def __init__(self, settings: Settings): self.settings = settings diff --git a/ai_search_agent/audit.py b/ai_search_agent/audit.py index 2c08d03..7136ac2 100644 --- a/ai_search_agent/audit.py +++ b/ai_search_agent/audit.py @@ -28,6 +28,9 @@ class AuditLogger: "usable_page_count": len([p for p in response.pages if p.usable]), "citation_count": len(response.answer.citations), "search_round_count": len(response.search_rounds), + "knowledge_triple_count": len(response.knowledge_triples), + "cache_hit": response.cache_hit, + "confidence": response.answer.confidence, }, "search_results": [ { @@ -58,6 +61,10 @@ class AuditLogger: } for item in response.search_rounds ], + "knowledge_triples": [ + {"subject": t.subject, "predicate": t.predicate, "object": t.object} + for t in response.knowledge_triples + ], } with self.log_path.open("a", encoding="utf-8") as fh: fh.write(json.dumps(record, ensure_ascii=False) + "\n") diff --git a/ai_search_agent/cache.py b/ai_search_agent/cache.py new file mode 100644 index 0000000..960e17d --- /dev/null +++ b/ai_search_agent/cache.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +import hashlib +import json + +import redis.asyncio as aioredis + +from .config import Settings + + +class SearchCache: + def __init__(self, settings: Settings): + self._url = settings.redis_url + self._ttl = settings.cache_ttl + self._client: aioredis.Redis | None = None + + async def _conn(self) -> aioredis.Redis | None: + if not self._url: + return None + if self._client is None: + self._client = aioredis.from_url( + self._url, + decode_responses=True, + socket_connect_timeout=5, + socket_timeout=5, + ) + return self._client + + @staticmethod + def _search_key(query: str, top_k: int, max_chars: int) -> str: + raw = f"{query}|{top_k}|{max_chars}" + h = hashlib.sha256(raw.encode()).hexdigest()[:20] + return f"aisou:search:{h}" + + async def get_search(self, query: str, top_k: int, max_chars: int) -> dict | None: + try: + c = await self._conn() + if c is None: + return None + raw = await c.get(self._search_key(query, top_k, max_chars)) + return json.loads(raw) if raw else None + except Exception: + return None + + async def set_search(self, query: str, top_k: int, max_chars: int, data: dict) -> None: + try: + c = await self._conn() + if c is None: + return + await c.setex( + self._search_key(query, top_k, max_chars), + self._ttl, + json.dumps(data, ensure_ascii=False), + ) + except Exception: + pass + + async def get_job(self, job_id: str) -> dict | None: + try: + c = await self._conn() + if c is None: + return None + raw = await c.get(f"aisou:job:{job_id}") + return json.loads(raw) if raw else None + except Exception: + return None + + async def set_job(self, job_id: str, data: dict, ttl: int = 7200) -> None: + try: + c = await self._conn() + if c is None: + return + await c.setex( + f"aisou:job:{job_id}", + ttl, + json.dumps(data, ensure_ascii=False), + ) + except Exception: + pass + + async def push_history(self, entry: dict) -> None: + try: + c = await self._conn() + if c is None: + return + await c.lpush("aisou:history", json.dumps(entry, ensure_ascii=False)) + await c.ltrim("aisou:history", 0, 99) + except Exception: + pass + + async def get_history(self, limit: int = 10) -> list[dict]: + try: + c = await self._conn() + if c is None: + return [] + items = await c.lrange("aisou:history", 0, limit - 1) + result: list[dict] = [] + for item in items: + try: + result.append(json.loads(item)) + except Exception: + pass + return result + except Exception: + return [] diff --git a/ai_search_agent/config.py b/ai_search_agent/config.py index 0aa670b..fb6322d 100644 --- a/ai_search_agent/config.py +++ b/ai_search_agent/config.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +from urllib.parse import quote from dotenv import load_dotenv from pydantic import BaseModel, Field @@ -24,6 +25,25 @@ class Settings(BaseModel): retry_attempts: int = Field(default=2, ge=0, le=5) max_search_rounds: int = Field(default=2, ge=1, le=5) audit_log_path: str = "logs/audit.jsonl" + # Redis cache + redis_url: str = "" + cache_ttl: int = Field(default=3600, ge=60) + # Feature flags + confidence_threshold: float = Field(default=0.6, ge=0.0, le=1.0) + enable_reranker: bool = True + enable_kg: bool = True + enable_memory: bool = True + enable_cache: bool = True + + @property + def llm_chat_endpoint(self) -> str: + """Always returns an OpenAI-compatible /chat/completions endpoint.""" + if self.llm_api_style == "openai_chat": + return self.llm_endpoint + base = self.llm_endpoint.rstrip("/") + if base.endswith("/responses"): + return base[: -len("/responses")] + "/chat/completions" + return base + "/chat/completions" @classmethod def from_env(cls) -> "Settings": @@ -59,6 +79,12 @@ class Settings(BaseModel): else: llm_endpoint = llm_endpoint or "https://api.x.ai/v1/responses" + redis_host = os.getenv("REDIS_HOST", "").strip() + redis_key = os.getenv("REDIS_KEY", "").strip() + redis_url = os.getenv("REDIS_URL", "").strip() + if not redis_url and redis_host and redis_key: + redis_url = f"rediss://:{quote(redis_key, safe='')}@{redis_host}:6380/0" + return cls( jina_api_key=jina_api_key, llm_api_key=llm_api_key, @@ -76,4 +102,11 @@ class Settings(BaseModel): retry_attempts=int(os.getenv("RETRY_ATTEMPTS", "2")), max_search_rounds=int(os.getenv("MAX_SEARCH_ROUNDS", "2")), audit_log_path=os.getenv("AUDIT_LOG_PATH", "logs/audit.jsonl").strip(), + redis_url=redis_url, + cache_ttl=int(os.getenv("CACHE_TTL", "3600")), + confidence_threshold=float(os.getenv("CONFIDENCE_THRESHOLD", "0.6")), + enable_reranker=os.getenv("ENABLE_RERANKER", "true").lower() == "true", + enable_kg=os.getenv("ENABLE_KG", "true").lower() == "true", + enable_memory=os.getenv("ENABLE_MEMORY", "true").lower() == "true", + enable_cache=os.getenv("ENABLE_CACHE", "true").lower() == "true", ) diff --git a/ai_search_agent/knowledge_graph.py b/ai_search_agent/knowledge_graph.py new file mode 100644 index 0000000..72591bd --- /dev/null +++ b/ai_search_agent/knowledge_graph.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import asyncio +import json + +import httpx + +from .config import Settings +from .models import KnowledgeTriple +from .parsers import extract_json_object +from .prompts import KG_EXTRACT_PROMPT + + +class KnowledgeGraphExtractor: + """Extract entity-relation-entity triples from page text using LLM.""" + + def __init__(self, settings: Settings): + self.settings = settings + + async def extract(self, text: str, query: str) -> list[KnowledgeTriple]: + if not text.strip(): + return [] + return await asyncio.to_thread(self._extract_sync, text[:3000], query) + + def _extract_sync(self, text: str, query: str) -> list[KnowledgeTriple]: + prompt = KG_EXTRACT_PROMPT.format(query=query, text=text) + payload = { + "model": self.settings.llm_model, + "messages": [{"role": "user", "content": prompt}], + "temperature": 0.0, + "max_tokens": 500, + } + try: + with httpx.Client(timeout=30) as client: + resp = client.post( + self.settings.llm_chat_endpoint, + headers={ + "Authorization": f"Bearer {self.settings.llm_api_key}", + "Content-Type": "application/json", + }, + json=payload, + ) + resp.raise_for_status() + text_out = resp.json()["choices"][0]["message"]["content"].strip() + raw = extract_json_object(text_out) + if raw: + triples = json.loads(raw).get("triples", []) + result: list[KnowledgeTriple] = [] + for t in triples[:5]: + try: + result.append(KnowledgeTriple(**t)) + except Exception: + pass + return result + except Exception: + pass + return [] diff --git a/ai_search_agent/llm_client.py b/ai_search_agent/llm_client.py index 36c1133..4c17dfb 100644 --- a/ai_search_agent/llm_client.py +++ b/ai_search_agent/llm_client.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +from collections.abc import AsyncGenerator import httpx @@ -90,3 +91,66 @@ class LLMClient: pass return AnswerPayload(summary=text.strip(), key_points=[], caveats=[], citations=[]) + + # ------------------------------------------------------------------ # + # Token-level streaming # + # ------------------------------------------------------------------ # + + async def answer_stream( + self, + system_prompt: str, + user_prompt: str, + allowed_urls: set[str], + ) -> AsyncGenerator[tuple[str, AnswerPayload | None], None]: + """Yield (token, None) for each streaming token, then ("", AnswerPayload) when done.""" + if self.settings.llm_api_style == "openai_chat": + full_text = "" + async for chunk in self._stream_openai_chat(system_prompt, user_prompt): + full_text += chunk + yield chunk, None + parsed = self._parse_answer(full_text) + else: + text = await self._call_xai_responses(system_prompt, user_prompt) + yield text, None + parsed = self._parse_answer(text) + + filtered = [c for c in parsed.citations if c.url in allowed_urls] + yield "", parsed.model_copy(update={"citations": filtered}) + + async def _stream_openai_chat( + self, system_prompt: str, user_prompt: str + ) -> AsyncGenerator[str, None]: + payload = { + "model": self.settings.llm_model, + "messages": [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ], + "temperature": 0.2, + "stream": True, + } + async with httpx.AsyncClient(timeout=self.settings.model_timeout) as client: + async with client.stream( + "POST", + self.settings.llm_endpoint, + headers={ + "Authorization": f"Bearer {self.settings.llm_api_key}", + "Content-Type": "application/json", + }, + json=payload, + ) as response: + response.raise_for_status() + async for line in response.aiter_lines(): + line = line.strip() + if not line.startswith("data: "): + continue + data_str = line[6:] + if data_str == "[DONE]": + break + try: + chunk_data = json.loads(data_str) + delta = chunk_data["choices"][0]["delta"].get("content", "") + if delta: + yield delta + except Exception: + pass diff --git a/ai_search_agent/models.py b/ai_search_agent/models.py index 2c9b406..5782de7 100644 --- a/ai_search_agent/models.py +++ b/ai_search_agent/models.py @@ -7,6 +7,7 @@ class SearchRequest(BaseModel): query: str = Field(min_length=1, description="User search query") top_k_pages: int | None = Field(default=None, ge=1, le=10) max_page_chars: int | None = Field(default=None, ge=500, le=20000) + callback_url: str | None = None class SearchResult(BaseModel): @@ -28,6 +29,13 @@ class PageContent(BaseModel): class Citation(BaseModel): title: str = "" url: str + trust_score: float = Field(default=0.6, ge=0.0, le=1.0) + + +class KnowledgeTriple(BaseModel): + subject: str + predicate: str + object: str class AnswerPayload(BaseModel): @@ -35,10 +43,13 @@ class AnswerPayload(BaseModel): key_points: list[str] = Field(default_factory=list) caveats: list[str] = Field(default_factory=list) citations: list[Citation] = Field(default_factory=list) + follow_up_questions: list[str] = Field(default_factory=list) + confidence: float = Field(default=0.8, ge=0.0, le=1.0) class SearchRoundTrace(BaseModel): round_index: int query: str + queries: list[str] = Field(default_factory=list) result_count: int = 0 fetched_page_count: int = 0 usable_page_count: int = 0 @@ -50,7 +61,9 @@ class SearchResponse(BaseModel): search_results: list[SearchResult] = Field(default_factory=list) pages: list[PageContent] = Field(default_factory=list) search_rounds: list[SearchRoundTrace] = Field(default_factory=list) + knowledge_triples: list[KnowledgeTriple] = Field(default_factory=list) content_ready: bool = True llm_called: bool = False + cache_hit: bool = False error: str | None = None audit_id: str | None = None diff --git a/ai_search_agent/parsers.py b/ai_search_agent/parsers.py index f1c478c..d7a274d 100644 --- a/ai_search_agent/parsers.py +++ b/ai_search_agent/parsers.py @@ -1,9 +1,42 @@ from __future__ import annotations import re +from urllib.parse import urlparse from .models import SearchResult +_DOMAIN_TRUST: dict[str, float] = { + "arxiv.org": 0.95, + "docs.python.org": 0.95, + "developer.mozilla.org": 0.95, + "pypi.org": 0.90, + "github.com": 0.90, + "wikipedia.org": 0.85, + "readthedocs.io": 0.85, + "microsoft.com": 0.85, + "azure.microsoft.com": 0.85, + "cloud.google.com": 0.85, + "aws.amazon.com": 0.85, + "stackoverflow.com": 0.80, + "npmjs.com": 0.80, + "medium.com": 0.55, + "zhihu.com": 0.55, + "csdn.net": 0.50, + "juejin.cn": 0.55, +} + + +def get_domain_trust(url: str) -> float: + """Return a 0-1 trust score based on the URL's domain.""" + try: + host = urlparse(url).hostname or "" + for domain, score in _DOMAIN_TRUST.items(): + if host == domain or host.endswith(f".{domain}"): + return score + except Exception: + pass + return 0.60 + MARKDOWN_LINK_RE = re.compile(r"\[(?P[^\]]+)\]\((?P<url>https?://[^\s)]+)\)") URL_RE = re.compile(r"https?://[^\s)>\]]+") diff --git a/ai_search_agent/planner.py b/ai_search_agent/planner.py new file mode 100644 index 0000000..891e98d --- /dev/null +++ b/ai_search_agent/planner.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import asyncio +import json +import re + +import httpx + +from .config import Settings +from .parsers import extract_json_object +from .prompts import ( + PLANNER_DECOMPOSE_PROMPT, + PLANNER_NEXT_QUERY_PROMPT, + PLANNER_TRANSLATE_PROMPT, +) + +_CJK_RE = re.compile(r"[\u4e00-\u9fff\u3040-\u30ff\uac00-\ud7af]") + + +class QueryPlanner: + """LLM-driven query decomposition, multilingual expansion, and next-round planning.""" + + def __init__(self, settings: Settings): + self.settings = settings + + async def decompose(self, query: str) -> list[str]: + prompt = PLANNER_DECOMPOSE_PROMPT.format(query=query) + result = await asyncio.to_thread(self._call_sync, prompt, max_tokens=300) + raw = extract_json_object(result) + if raw: + try: + subs = json.loads(raw).get("sub_queries", []) + cleaned = [s.strip() for s in subs if s.strip()] + if cleaned: + return cleaned + except Exception: + pass + return [query] + + async def expand_multilingual(self, query: str) -> list[str]: + if not _CJK_RE.search(query): + return [query] + prompt = PLANNER_TRANSLATE_PROMPT.format(query=query) + result = await asyncio.to_thread(self._call_sync, prompt, max_tokens=100) + raw = extract_json_object(result) + if raw: + try: + en = json.loads(raw).get("en", "").strip() + if en and en.lower() != query.lower(): + return [query, en] + except Exception: + pass + return [query] + + async def next_query( + self, + original_query: str, + round_index: int, + search_digest: str, + usable_count: int, + ) -> str: + prompt = PLANNER_NEXT_QUERY_PROMPT.format( + query=original_query, + round_index=round_index, + usable_count=usable_count, + digest=search_digest[:600], + ) + result = await asyncio.to_thread(self._call_sync, prompt, max_tokens=150) + raw = extract_json_object(result) + if raw: + try: + return json.loads(raw).get("next_query", "").strip() + except Exception: + pass + return "" + + def _call_sync(self, user_prompt: str, max_tokens: int = 200) -> str: + payload = { + "model": self.settings.llm_model, + "messages": [{"role": "user", "content": user_prompt}], + "temperature": 0.1, + "max_tokens": max_tokens, + } + try: + with httpx.Client(timeout=30) as client: + resp = client.post( + self.settings.llm_chat_endpoint, + headers={ + "Authorization": f"Bearer {self.settings.llm_api_key}", + "Content-Type": "application/json", + }, + json=payload, + ) + resp.raise_for_status() + return resp.json()["choices"][0]["message"]["content"].strip() + except Exception: + return "" diff --git a/ai_search_agent/prompts.py b/ai_search_agent/prompts.py index fac3bad..d9936fe 100644 --- a/ai_search_agent/prompts.py +++ b/ai_search_agent/prompts.py @@ -1,15 +1,43 @@ SYSTEM_PROMPT = """You are an evidence-grounded web research assistant. -Return valid JSON only with this shape: +Return valid JSON only with this exact shape (no extra keys, no markdown fences): { - "summary": "short answer in the user's language", + "summary": "concise answer in the user's language", "key_points": ["point 1", "point 2"], "caveats": ["optional limitation"], - "citations": [{"title": "source title", "url": "https://..."}] + "citations": [{"title": "source title", "url": "https://..."}], + "follow_up_questions": ["related question 1", "related question 2", "related question 3"], + "confidence": 0.85 } Rules: - Use only the supplied search digest and page excerpts. -- Do not invent citations. +- Do not invent citations; only cite URLs from the fetched pages. - Prefer concise, factual writing. - If the evidence is incomplete or conflicting, mention it in caveats. +- confidence: float 0.0-1.0, how confident you are the answer is complete and accurate. +- follow_up_questions: exactly 3 relevant follow-up questions in the user's language. """ + +PLANNER_DECOMPOSE_PROMPT = """Decompose the following search query into 1-3 focused sub-queries for web search. +Return JSON only: {{"sub_queries": ["...", "..."]}} +If no decomposition is needed, return just the original query in the list. +Query: {query}""" + +PLANNER_NEXT_QUERY_PROMPT = """You are a search planner. +Original query: {query} +Current round: {round_index} +Usable pages found: {usable_count} +Search digest (first 600 chars): {digest} + +If more evidence is needed, generate exactly one better follow-up search query in the same language as the original query. +Return JSON only: {{"next_query": "..."}} +If evidence is sufficient, return {{"next_query": ""}}""" + +PLANNER_TRANSLATE_PROMPT = """Translate this search query to English. +Return JSON only: {{"en": "..."}} +Query: {query}""" + +KG_EXTRACT_PROMPT = """Extract up to 5 key entity-relation-entity triples from the text that are relevant to the query. +Return JSON only: {{"triples": [{{"subject": "...", "predicate": "...", "object": "..."}}]}} +Query: {query} +Text: {text}""" diff --git a/ai_search_agent/reranker.py b/ai_search_agent/reranker.py new file mode 100644 index 0000000..02f381d --- /dev/null +++ b/ai_search_agent/reranker.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import asyncio + +import httpx + +from .config import Settings +from .models import PageContent + + +class JinaReranker: + """Rerank pages by semantic relevance using Jina Reranker v2.""" + + _ENDPOINT = "https://api.jina.ai/v1/rerank" + _MODEL = "jina-reranker-v2-base-multilingual" + + def __init__(self, settings: Settings): + self.settings = settings + + async def rerank(self, query: str, pages: list[PageContent]) -> list[PageContent]: + if len(pages) <= 1: + return pages + return await asyncio.to_thread(self._rerank_sync, query, pages) + + def _rerank_sync(self, query: str, pages: list[PageContent]) -> list[PageContent]: + documents = [ + {"text": (page.content[:2000] if page.content else page.title) or page.url} + for page in pages + ] + payload = { + "model": self._MODEL, + "query": query, + "documents": documents, + "top_n": len(pages), + } + try: + with httpx.Client(timeout=30) as client: + resp = client.post( + self._ENDPOINT, + headers={ + "Authorization": f"Bearer {self.settings.jina_api_key}", + "Content-Type": "application/json", + }, + json=payload, + ) + resp.raise_for_status() + results = resp.json().get("results", []) + return [pages[item["index"]] for item in results if item["index"] < len(pages)] + except Exception: + return pages # non-fatal: return original order diff --git a/requirements.txt b/requirements.txt index 7e3fa6a..a6c8c84 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ requests>=2.32.0 pydantic>=2.10.0 python-dotenv>=1.0.1 uvicorn>=0.34.0 +redis[asyncio]>=5.2.0 azure-functions>=1.21.0