commit 9ed6397dbc699d9e4ac5115d958c0ea21f7da243 Author: xiaohei Date: Fri Mar 27 17:06:54 2026 +0800 Initial AI search agent diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..5096c4c --- /dev/null +++ b/.env.example @@ -0,0 +1,14 @@ +JINA_API_KEY= +LLM_API_KEY= +LLM_API_STYLE=openai_chat +LLM_ENDPOINT=https://dashscope-intl.aliyuncs.com/compatible-mode/v1 +MODEL_NAME=qwen3.5-flash +SEARCH_PROVIDER=bing_html +TOP_K_PAGES=5 +MAX_PAGE_CHARS=5000 +MAX_CONCURRENCY=5 +SEARCH_TIMEOUT=20 +READER_TIMEOUT=25 +MODEL_TIMEOUT=90 +RETRY_ATTEMPTS=2 +AUDIT_LOG_PATH=logs/audit.jsonl diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2486a9d --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +.venv/ +__pycache__/ +*.pyc +logs/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..8284e2c --- /dev/null +++ b/README.md @@ -0,0 +1,127 @@ +# AI Search Agent + +独立实现的 AI 搜索 Agent，不复用仓库里原有的 `search_agent`。 +这个版本按“较完整的最佳实践骨架”组织，不是最小脚本： + +1. 用 `Bing HTML` 做网页搜索 +2. 用 `r.jina.ai` 并发深读前几条结果 +3. 用可切换的 LLM 生成结构化回答 +4. 暴露 CLI 和 FastAPI 两种入口 + +## 目录 + +```text +ai_search_agent/ +├── ai_search_agent/ +│ ├── agent.py +│ ├── api.py +│ ├── config.py +│ ├── jina.py +│ ├── llm_client.py +│ ├── models.py +│ ├── parsers.py +│ ├── prompts.py +│ └── __init__.py +├── .env.example +├── tests/ +│ └── test_parsers.py +├── main.py +├── README.md +├── serve.py +└── requirements.txt +``` + +## 环境变量 + +```bash +export JINA_API_KEY=your-jina-key +export LLM_API_KEY=your-llm-key +export LLM_API_STYLE=openai_chat +export LLM_ENDPOINT=https://dashscope-intl.aliyuncs.com/compatible-mode/v1 +export MODEL_NAME=qwen3.5-flash +export SEARCH_PROVIDER=bing_html +export TOP_K_PAGES=5 +export MAX_PAGE_CHARS=5000 +export MAX_CONCURRENCY=5 +export SEARCH_TIMEOUT=20 +export READER_TIMEOUT=25 +export MODEL_TIMEOUT=90 +export RETRY_ATTEMPTS=2 +export AUDIT_LOG_PATH=logs/audit.jsonl +``` + +如果切回 xAI： + +```bash +export JINA_API_KEY=your-jina-key +export LLM_API_KEY=your-xai-key +export LLM_API_STYLE=xai_responses +export MODEL_NAME=grok-4.20-multi-agent-0309 +``` + +## 安装 + +```bash +cd ai_search_agent +python3 -m venv .venv +source .venv/bin/activate +pip install -r requirements.txt +``` + +## 运行 + +```bash +python main.py "PydanticAI 适合哪些场景？" +``` + +输出 JSON： + +```bash +python main.py "PydanticAI 适合哪些场景？" --json +``` + +## 启动 API 服务 + +```bash +uvicorn serve:app --host 0.0.0.0 --port 8080 +``` + +请求示例： + +```bash +curl -X POST http://127.0.0.1:8080/search \ + -H "Content-Type: application/json" \ + -d '{"query":"PydanticAI 适合哪些场景？","top_k_pages":4}' +``` + +## 设计说明 + +- `SearchClient.search()` 固定走 `bing_html` +- `JinaClient.read_pages()` 只负责通过 `r.jina.ai` 并发抓取页面内容，并带重试 +- `LLMClient.answer()` 支持 `xai_responses` 和 `openai_chat` 两种模式 +- 每次请求都会生成 `audit_id`，并把审计记录落到 `logs/audit.jsonl` +- 模型输出约束为 JSON 形状，解析失败时退化为纯文本答案 +- `AISearchAgent.run()` 统一编排搜索、抓取、归纳和来源回填 +- `create_app()` 复用同一套业务链路，不重复实现 API 逻辑 + +## 返回结构 + +```json +{ + "query": "PydanticAI 适合哪些场景？", + "answer": { + "summary": "...", + "key_points": ["..."], + "caveats": ["..."], + "citations": [{"title": "...", "url": "https://..."}] + }, + "search_results": [], + "pages": [] +} +``` + +## 当前边界 + +- 还没做 SSE 流式输出 +- 还没做多轮搜索规划和反思 +- Bing HTML 解析依赖当前结果页 DOM 结构 diff --git a/ai_search_agent/__init__.py b/ai_search_agent/__init__.py new file mode 100644 index 0000000..72a4851 --- /dev/null +++ b/ai_search_agent/__init__.py @@ -0,0 +1 @@ +"""Standalone AI search agent built on Jina + xAI.""" diff --git a/ai_search_agent/agent.py b/ai_search_agent/agent.py new file mode 100644 index 0000000..2ce2492 --- /dev/null +++ b/ai_search_agent/agent.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from .config import Settings +from .jina import JinaClient +from .llm_client import LLMClient +from .models import AnswerPayload, Citation, PageContent, SearchRequest, SearchResponse, SearchResult +from .prompts import SYSTEM_PROMPT +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) + + async def run(self, request: SearchRequest) -> SearchResponse: + top_k_pages = request.top_k_pages or self.settings.top_k_pages + max_page_chars = request.max_page_chars or self.settings.max_page_chars + + raw_search_text, search_results = await self.search.search(request.query) + selected_results = search_results[: max(top_k_pages * 3, top_k_pages)] + pages = await self.jina.read_pages(selected_results, max_page_chars=max_page_chars) + usable_pages = [page for page in pages if page.fetched and page.usable] + + if not usable_pages: + return SearchResponse( + query=request.query, + answer=AnswerPayload( + summary="未读取到可用正文内容，已拒绝调用模型总结。", + key_points=[], + caveats=[ + "所有候选页面都未能提取到有效正文，可能是 403、451、反爬或页面过短。", + "请更换查询词，或增加可抓取来源。", + ], + citations=[], + ), + search_results=selected_results, + pages=pages, + content_ready=False, + llm_called=False, + error="NO_USABLE_CONTENT", + ) + + pages_for_model = usable_pages[:top_k_pages] + results_for_model = [ + result for result in selected_results if any(page.url == result.url for page in pages_for_model) + ] + answer = await self.llm.answer( + SYSTEM_PROMPT, + self._build_user_prompt(request.query, raw_search_text, results_for_model, pages_for_model), + allowed_urls={page.url for page in pages_for_model if page.fetched}, + ) + + if not answer.citations: + answer = self._backfill_citations(answer, results_for_model) + + return SearchResponse( + query=request.query, + answer=answer, + search_results=results_for_model, + pages=pages, + content_ready=True, + llm_called=True, + ) + + def _build_user_prompt( + self, + query: str, + raw_search_text: str, + search_results: list[SearchResult], + pages: list[PageContent], + ) -> str: + result_lines: list[str] = [] + for result in search_results: + result_lines.append(f"[Result {result.rank}] {result.title or 'Untitled'}") + result_lines.append(f"URL: {result.url}") + if result.snippet: + result_lines.append(f"Snippet: {result.snippet}") + + page_lines: list[str] = [] + for index, page in enumerate(pages, start=1): + page_lines.append(f"[Page {index}] {page.title or page.url}") + page_lines.append(f"URL: {page.url}") + if page.fetched: + page_lines.append(page.content) + else: + page_lines.append(f"Fetch error: {page.error}") + + return f"""User query: +{query} + +Search digest: +{raw_search_text} + +Selected search results: +{chr(10).join(result_lines)} + +Fetched page excerpts: +{chr(10).join(page_lines)} +""" + + def _backfill_citations( + self, + answer: AnswerPayload, + search_results: list[SearchResult], + ) -> AnswerPayload: + fallback = [ + Citation(title=item.title or item.url, url=item.url) + for item in search_results[:3] + ] + return answer.model_copy(update={"citations": fallback}) diff --git a/ai_search_agent/api.py b/ai_search_agent/api.py new file mode 100644 index 0000000..ebc81da --- /dev/null +++ b/ai_search_agent/api.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from fastapi import FastAPI +from fastapi.responses import JSONResponse + +from .agent import AISearchAgent +from .audit import AuditLogger +from .config import Settings +from .models import SearchRequest + + +def create_app() -> FastAPI: + settings = Settings.from_env() + agent = AISearchAgent(settings) + audit_logger = AuditLogger(settings.audit_log_path) + app = FastAPI(title="AI Search Agent", version="1.0.0") + + @app.get("/health") + async def health() -> dict[str, str]: + return {"status": "ok"} + + @app.post("/search") + async def search(request: SearchRequest) -> JSONResponse: + response = await agent.run(request) + audit_id = audit_logger.write(request, response) + response = response.model_copy(update={"audit_id": audit_id}) + return JSONResponse(content=response.model_dump()) + + return app diff --git a/ai_search_agent/audit.py b/ai_search_agent/audit.py new file mode 100644 index 0000000..9baecb9 --- /dev/null +++ b/ai_search_agent/audit.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path +from uuid import uuid4 + +from .models import SearchRequest, SearchResponse + + +class AuditLogger: + def __init__(self, log_path: str | Path): + self.log_path = Path(log_path) + self.log_path.parent.mkdir(parents=True, exist_ok=True) + + def write(self, request: SearchRequest, response: SearchResponse) -> str: + audit_id = uuid4().hex + record = { + "audit_id": audit_id, + "timestamp": datetime.now(timezone.utc).isoformat(), + "request": request.model_dump(), + "response_summary": { + "content_ready": response.content_ready, + "llm_called": response.llm_called, + "error": response.error, + "search_result_count": len(response.search_results), + "page_count": len(response.pages), + "usable_page_count": len([p for p in response.pages if p.usable]), + "citation_count": len(response.answer.citations), + }, + "search_results": [ + { + "rank": item.rank, + "title": item.title, + "url": item.url, + } + for item in response.search_results + ], + "pages": [ + { + "url": item.url, + "title": item.title, + "fetched": item.fetched, + "usable": item.usable, + "error": item.error, + } + for item in response.pages + ], + } + with self.log_path.open("a", encoding="utf-8") as fh: + fh.write(json.dumps(record, ensure_ascii=False) + "\n") + return audit_id diff --git a/ai_search_agent/config.py b/ai_search_agent/config.py new file mode 100644 index 0000000..1df34eb --- /dev/null +++ b/ai_search_agent/config.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import os + +from dotenv import load_dotenv +from pydantic import BaseModel, Field + + +class Settings(BaseModel): + jina_api_key: str + llm_api_key: str + llm_model: str = "qwen3.5-flash" + llm_api_style: str = "openai_chat" + search_provider: str = "bing_html" + bing_search_url: str = "https://www.bing.com/search" + reader_base_url: str = "https://r.jina.ai/" + llm_endpoint: str = "https://dashscope-intl.aliyuncs.com/compatible-mode/v1/chat/completions" + search_timeout: float = 20.0 + reader_timeout: float = 25.0 + model_timeout: float = 90.0 + top_k_pages: int = Field(default=5, ge=1, le=10) + max_page_chars: int = Field(default=5000, ge=500, le=20000) + max_concurrency: int = Field(default=5, ge=1, le=20) + retry_attempts: int = Field(default=2, ge=0, le=5) + audit_log_path: str = "logs/audit.jsonl" + + @classmethod + def from_env(cls) -> "Settings": + load_dotenv() + jina_api_key = os.getenv("JINA_API_KEY", "").strip() + llm_api_key = ( + os.getenv("LLM_API_KEY", "").strip() + or os.getenv("OPENAI_API_KEY", "").strip() + or os.getenv("XAI_API_KEY", "").strip() + ) + + if not jina_api_key: + raise ValueError("Missing JINA_API_KEY") + if not llm_api_key: + raise ValueError("Missing LLM_API_KEY") + + llm_api_style = os.getenv("LLM_API_STYLE", "openai_chat").strip() + llm_model = ( + os.getenv("MODEL_NAME", "").strip() + or os.getenv("LLM_MODEL", "").strip() + or os.getenv("XAI_MODEL", "qwen3.5-flash").strip() + ) + llm_endpoint = ( + os.getenv("LLM_ENDPOINT", "").strip() + or os.getenv("OPENAI_BASE_URL", "").strip() + or os.getenv("LLM_BASE_URL", "").strip() + ) + + if llm_api_style == "openai_chat": + llm_endpoint = llm_endpoint or "https://dashscope-intl.aliyuncs.com/compatible-mode/v1" + if not llm_endpoint.endswith("/chat/completions"): + llm_endpoint = llm_endpoint.rstrip("/") + "/chat/completions" + else: + llm_endpoint = llm_endpoint or "https://api.x.ai/v1/responses" + + return cls( + jina_api_key=jina_api_key, + llm_api_key=llm_api_key, + llm_model=llm_model, + llm_api_style=llm_api_style, + search_provider=os.getenv("SEARCH_PROVIDER", "bing_html").strip(), + llm_endpoint=llm_endpoint, + search_timeout=float(os.getenv("SEARCH_TIMEOUT", "20")), + reader_timeout=float(os.getenv("READER_TIMEOUT", "25")), + model_timeout=float(os.getenv("MODEL_TIMEOUT", "90")), + top_k_pages=int(os.getenv("TOP_K_PAGES", "5")), + max_page_chars=int(os.getenv("MAX_PAGE_CHARS", "5000")), + max_concurrency=int(os.getenv("MAX_CONCURRENCY", "5")), + retry_attempts=int(os.getenv("RETRY_ATTEMPTS", "2")), + audit_log_path=os.getenv("AUDIT_LOG_PATH", "logs/audit.jsonl").strip(), + ) diff --git a/ai_search_agent/jina.py b/ai_search_agent/jina.py new file mode 100644 index 0000000..e25f35c --- /dev/null +++ b/ai_search_agent/jina.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import asyncio + +import httpx + +from .config import Settings +from .models import PageContent, SearchResult +from .parsers import trim_text + + +class JinaClient: + def __init__(self, settings: Settings): + self.settings = settings + self._semaphore = asyncio.Semaphore(settings.max_concurrency) + + @property + def _headers(self) -> dict[str, str]: + return { + "Authorization": f"Bearer {self.settings.jina_api_key}", + "X-Respond-With": "readerlm-v2", + "X-Retain-Images": "none", + "X-Retain-Links": "gpt-oss", + } + + async def read_pages(self, results: list[SearchResult], max_page_chars: int) -> list[PageContent]: + tasks = [self.read_page(result.url, result.title, max_page_chars) for result in results] + return await asyncio.gather(*tasks) + + async def read_page(self, url: str, fallback_title: str, max_page_chars: int) -> PageContent: + reader_url = f"{self.settings.reader_base_url}{url}" + try: + text = await self._get_text( + reader_url, + timeout=self.settings.reader_timeout, + extra_headers={"Accept": "text/plain"}, + use_reader_headers=False, + ) + title = "" + first_line = text.splitlines()[0].strip() if text else "" + if first_line.startswith("#"): + title = first_line.lstrip("#").strip() + + return PageContent( + url=url, + title=title or fallback_title, + content=trim_text(text.strip(), max_page_chars), + fetched=True, + usable=self._is_usable_content(text), + ) + except Exception as exc: + return PageContent( + url=url, + title=fallback_title, + content="", + fetched=False, + usable=False, + error=str(exc), + ) + + def _is_usable_content(self, text: str) -> bool: + normalized = text.strip() + if len(normalized) < 300: + return False + + blocked_markers = [ + "403: Forbidden", + "451 Unavailable For Legal Reasons", + "Target URL returned error 403", + "Target URL returned error 451", + "enable javascript", + "captcha", + ] + lowered = normalized.lower() + return not any(marker.lower() in lowered for marker in blocked_markers) + + async def _get_text( + self, + url: str, + timeout: float, + extra_headers: dict[str, str] | None = None, + use_reader_headers: bool = True, + ) -> str: + headers = dict(self._headers) if use_reader_headers else {} + if extra_headers: + headers.update(extra_headers) + + last_error: Exception | None = None + for _ in range(self.settings.retry_attempts + 1): + try: + async with self._semaphore: + async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client: + response = await client.get(url, headers=headers) + response.raise_for_status() + return response.text + except Exception as exc: + last_error = exc + await asyncio.sleep(0.5) + + assert last_error is not None + raise last_error diff --git a/ai_search_agent/llm_client.py b/ai_search_agent/llm_client.py new file mode 100644 index 0000000..36c1133 --- /dev/null +++ b/ai_search_agent/llm_client.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import asyncio +import json + +import httpx + +from .config import Settings +from .models import AnswerPayload +from .parsers import extract_json_object + + +class LLMClient: + def __init__(self, settings: Settings): + self.settings = settings + + async def answer(self, system_prompt: str, user_prompt: str, allowed_urls: set[str]) -> AnswerPayload: + if self.settings.llm_api_style == "openai_chat": + text = await self._call_openai_chat(system_prompt, user_prompt) + else: + text = await self._call_xai_responses(system_prompt, user_prompt) + + parsed = self._parse_answer(text) + filtered_citations = [item for item in parsed.citations if item.url in allowed_urls] + return parsed.model_copy(update={"citations": filtered_citations}) + + async def _call_xai_responses(self, system_prompt: str, user_prompt: str) -> str: + return await asyncio.to_thread(self._call_xai_responses_sync, system_prompt, user_prompt) + + def _call_xai_responses_sync(self, system_prompt: str, user_prompt: str) -> str: + payload = { + "model": self.settings.llm_model, + "input": [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ], + } + with httpx.Client(timeout=self.settings.model_timeout) as client: + response = client.post( + self.settings.llm_endpoint, + headers={ + "Authorization": f"Bearer {self.settings.llm_api_key}", + "Content-Type": "application/json", + }, + json=payload, + ) + response.raise_for_status() + + data = response.json() + output_text: list[str] = [] + for item in data.get("output", []): + for content in item.get("content", []): + if content.get("type") == "output_text": + output_text.append(content.get("text", "")) + return "\n".join(part for part in output_text if part).strip() + + async def _call_openai_chat(self, system_prompt: str, user_prompt: str) -> str: + return await asyncio.to_thread(self._call_openai_chat_sync, system_prompt, user_prompt) + + def _call_openai_chat_sync(self, system_prompt: str, user_prompt: str) -> str: + payload = { + "model": self.settings.llm_model, + "messages": [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ], + "temperature": 0.2, + } + with httpx.Client(timeout=self.settings.model_timeout) as client: + response = client.post( + self.settings.llm_endpoint, + headers={ + "Authorization": f"Bearer {self.settings.llm_api_key}", + "Content-Type": "application/json", + }, + json=payload, + ) + response.raise_for_status() + + data = response.json() + return data["choices"][0]["message"]["content"].strip() + + def _parse_answer(self, text: str) -> AnswerPayload: + raw_json = extract_json_object(text) + if raw_json: + try: + data = json.loads(raw_json) + return AnswerPayload.model_validate(data) + except Exception: + pass + + return AnswerPayload(summary=text.strip(), key_points=[], caveats=[], citations=[]) diff --git a/ai_search_agent/models.py b/ai_search_agent/models.py new file mode 100644 index 0000000..8db310e --- /dev/null +++ b/ai_search_agent/models.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +from pydantic import BaseModel, Field + + +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) + + +class SearchResult(BaseModel): + url: str + title: str = "" + snippet: str = "" + rank: int = 0 + + +class PageContent(BaseModel): + url: str + title: str = "" + content: str = "" + fetched: bool = True + usable: bool = True + error: str | None = None + + +class Citation(BaseModel): + title: str = "" + url: str + + +class AnswerPayload(BaseModel): + summary: str + key_points: list[str] = Field(default_factory=list) + caveats: list[str] = Field(default_factory=list) + citations: list[Citation] = Field(default_factory=list) + + +class SearchResponse(BaseModel): + query: str + answer: AnswerPayload + search_results: list[SearchResult] = Field(default_factory=list) + pages: list[PageContent] = Field(default_factory=list) + content_ready: bool = True + llm_called: 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 new file mode 100644 index 0000000..f1c478c --- /dev/null +++ b/ai_search_agent/parsers.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +import re + +from .models import SearchResult + +MARKDOWN_LINK_RE = re.compile(r"\[(?P[^\]]+)\]\((?P<url>https?://[^\s)]+)\)") +URL_RE = re.compile(r"https?://[^\s)>\]]+") + + +def extract_markdown_links(text: str) -> list[SearchResult]: + results: list[SearchResult] = [] + seen: set[str] = set() + + for rank, match in enumerate(MARKDOWN_LINK_RE.finditer(text), start=1): + url = match.group("url").strip() + if url in seen: + continue + seen.add(url) + start = max(0, match.start() - 160) + end = min(len(text), match.end() + 160) + snippet = " ".join(text[start:end].split()) + results.append( + SearchResult( + title=match.group("title").strip(), + url=url, + snippet=snippet, + rank=rank, + ) + ) + + return results + + +def extract_urls(text: str) -> list[str]: + urls: list[str] = [] + seen: set[str] = set() + + for match in URL_RE.finditer(text): + url = match.group(0).rstrip(".,") + if url in seen: + continue + seen.add(url) + urls.append(url) + + return urls + + +def trim_text(text: str, limit: int) -> str: + if len(text) <= limit: + return text + return text[:limit].rstrip() + "\n...[truncated]" + + +def extract_json_object(text: str) -> str | None: + start = text.find("{") + end = text.rfind("}") + if start == -1 or end == -1 or end <= start: + return None + return text[start : end + 1] diff --git a/ai_search_agent/prompts.py b/ai_search_agent/prompts.py new file mode 100644 index 0000000..fac3bad --- /dev/null +++ b/ai_search_agent/prompts.py @@ -0,0 +1,15 @@ +SYSTEM_PROMPT = """You are an evidence-grounded web research assistant. +Return valid JSON only with this shape: +{ + "summary": "short answer in the user's language", + "key_points": ["point 1", "point 2"], + "caveats": ["optional limitation"], + "citations": [{"title": "source title", "url": "https://..."}] +} + +Rules: +- Use only the supplied search digest and page excerpts. +- Do not invent citations. +- Prefer concise, factual writing. +- If the evidence is incomplete or conflicting, mention it in caveats. +""" diff --git a/ai_search_agent/search_client.py b/ai_search_agent/search_client.py new file mode 100644 index 0000000..5ad48c2 --- /dev/null +++ b/ai_search_agent/search_client.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from urllib.parse import quote + +import httpx +from bs4 import BeautifulSoup + +from .config import Settings +from .models import SearchResult +class SearchClient: + def __init__(self, settings: Settings): + self.settings = settings + + async def search(self, query: str) -> tuple[str, list[SearchResult]]: + return await self._bing_search(query) + + async def _bing_search(self, query: str) -> tuple[str, list[SearchResult]]: + url = f"{self.settings.bing_search_url}?q={quote(query)}" + async with httpx.AsyncClient(timeout=self.settings.search_timeout, follow_redirects=True) as client: + response = await client.get( + url, + headers={ + "User-Agent": "Mozilla/5.0", + "Accept-Language": "en-US,en;q=0.9", + }, + ) + response.raise_for_status() + + html = response.text + soup = BeautifulSoup(html, "html.parser") + results: list[SearchResult] = [] + for rank, item in enumerate(soup.select("li.b_algo"), start=1): + link = item.select_one("h2 a") + if not link: + continue + href = (link.get("href") or "").strip() + if not href.startswith("http"): + continue + snippet_node = item.select_one(".b_caption p") + results.append( + SearchResult( + title=link.get_text(" ", strip=True), + url=href, + snippet=snippet_node.get_text(" ", strip=True) if snippet_node else "", + rank=rank, + ) + ) + if len(results) >= self.settings.top_k_pages * 2: + break + + if not results: + raise ValueError("No Bing HTML search results parsed") + return html, results diff --git a/main.py b/main.py new file mode 100644 index 0000000..90d9546 --- /dev/null +++ b/main.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import argparse +import asyncio +import json + +from ai_search_agent.agent import AISearchAgent +from ai_search_agent.config import Settings +from ai_search_agent.models import SearchRequest + + +async def _run_query(args: argparse.Namespace) -> None: + settings = Settings.from_env() + agent = AISearchAgent(settings) + result = await agent.run( + SearchRequest( + query=args.query, + top_k_pages=args.top_k_pages, + max_page_chars=args.max_page_chars, + ) + ) + + if args.json: + print(json.dumps(result.model_dump(), ensure_ascii=False, indent=2)) + return + + print(result.answer.summary) + if result.answer.key_points: + print("\nKey Points:") + for item in result.answer.key_points: + print(f"- {item}") + if result.answer.caveats: + print("\nCaveats:") + for item in result.answer.caveats: + print(f"- {item}") + if result.answer.citations: + print("\nSources:") + for item in result.answer.citations: + print(f"- {item.title}: {item.url}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Standalone AI search agent") + parser.add_argument("query", help="Search query") + parser.add_argument("--json", action="store_true", help="Print full JSON output") + parser.add_argument("--top-k-pages", type=int, default=None, help="How many pages to fetch") + parser.add_argument("--max-page-chars", type=int, default=None, help="Max chars per page") + args = parser.parse_args() + asyncio.run(_run_query(args)) + + +if __name__ == "__main__": + main() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..40667ee --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +fastapi>=0.115.0 +httpx>=0.28.0 +pydantic>=2.10.0 +python-dotenv>=1.0.1 +uvicorn>=0.34.0 +beautifulsoup4>=4.12.0 diff --git a/serve.py b/serve.py new file mode 100644 index 0000000..7cea5db --- /dev/null +++ b/serve.py @@ -0,0 +1,3 @@ +from ai_search_agent.api import create_app + +app = create_app() diff --git a/tests/test_parsers.py b/tests/test_parsers.py new file mode 100644 index 0000000..ae5612c --- /dev/null +++ b/tests/test_parsers.py @@ -0,0 +1,37 @@ +import unittest + +from ai_search_agent.parsers import extract_json_object, extract_markdown_links, extract_urls, trim_text + + +class ParserTests(unittest.TestCase): + def test_extract_markdown_links_deduplicates(self) -> None: + text = """ + [Example](https://example.com) + [Example Again](https://example.com) + [Docs](https://docs.example.com/path) + """ + results = extract_markdown_links(text) + + self.assertEqual( + [item.url for item in results], + ["https://example.com", "https://docs.example.com/path"], + ) + + def test_extract_urls(self) -> None: + text = "See https://a.example.com, https://b.example.com/docs." + self.assertEqual( + extract_urls(text), + ["https://a.example.com", "https://b.example.com/docs"], + ) + + def test_trim_text(self) -> None: + self.assertEqual(trim_text("abc", 5), "abc") + self.assertTrue(trim_text("abcdef", 3).startswith("abc")) + + def test_extract_json_object(self) -> None: + text = "prefix {\"summary\": \"ok\"} suffix" + self.assertEqual(extract_json_object(text), "{\"summary\": \"ok\"}") + + +if __name__ == "__main__": + unittest.main()