feat: 12项能力全实现 - Redis缓存/多轮LLM规划/多语言扩展/Jina重排/KG提取/置信度重搜/Token流式/引用可信度/追问建议/异步Job/Webhook/历史记忆

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

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")

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 []

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",
         )

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 []

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

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

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<title>[^\]]+)\]\((?P<url>https?://[^\s)]+)\)")
 URL_RE = re.compile(r"https?://[^\s)>\]]+")
 

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 ""

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}"""

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

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