wehub-resource-sync
263 lines
8.5 KiB
Python
263 lines
8.5 KiB
Python
#!/usr/bin/env python3
|
|
"""PixelRAG Agent — Claude + tool_use for visual web search.
|
|
|
|
A real Anthropic agent that uses Claude to answer questions by searching
|
|
a visual Wikipedia index via PixelRAG. Claude decides when to call the
|
|
search tool and synthesizes answers from visual retrieval results.
|
|
|
|
Prerequisites:
|
|
- ANTHROPIC_API_KEY env var set
|
|
- pixelrag serve running on localhost:30001 (or set --endpoint)
|
|
|
|
Usage:
|
|
# Interactive conversation with the agent
|
|
python demos/agent_skill.py
|
|
|
|
# Single question
|
|
python demos/agent_skill.py "Who invented the telephone?"
|
|
|
|
# Custom endpoint
|
|
python demos/agent_skill.py --endpoint http://gpu-box:30001 "Eiffel Tower history"
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import urllib.request
|
|
|
|
import anthropic
|
|
|
|
SEARCH_TOOL = {
|
|
"name": "pixelrag_search",
|
|
"description": (
|
|
"Search a visual Wikipedia index using natural language queries. "
|
|
"Returns ranked results with article URLs and relevance scores. "
|
|
"Use this tool to find information about any topic — it searches "
|
|
"screenshot-based embeddings of Wikipedia articles, so it works well "
|
|
"for both textual and visual content."
|
|
),
|
|
"input_schema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "Natural language search query",
|
|
},
|
|
"n_results": {
|
|
"type": "integer",
|
|
"description": "Number of results to return (default 5, max 20)",
|
|
"default": 5,
|
|
},
|
|
},
|
|
"required": ["query"],
|
|
},
|
|
}
|
|
|
|
WEB_FETCH_TOOL = {
|
|
"name": "web_fetch",
|
|
"description": (
|
|
"Fetch the text content of a URL. Use this to read Wikipedia articles "
|
|
"or other web pages found via search. Returns the page text content."
|
|
),
|
|
"input_schema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"url": {
|
|
"type": "string",
|
|
"description": "URL to fetch",
|
|
},
|
|
},
|
|
"required": ["url"],
|
|
},
|
|
}
|
|
|
|
TOOLS = [SEARCH_TOOL, WEB_FETCH_TOOL]
|
|
|
|
SYSTEM_PROMPT = """\
|
|
You are a research assistant with access to a visual Wikipedia search engine (PixelRAG).
|
|
When asked a question, use the pixelrag_search tool to find relevant Wikipedia articles,
|
|
then synthesize an answer from the results. You may search multiple times with different
|
|
queries to gather comprehensive information. Cite your sources with Wikipedia URLs.
|
|
|
|
If search results are insufficient, say so honestly rather than guessing."""
|
|
|
|
|
|
def execute_pixelrag_search(
|
|
query: str, n_results: int = 5, endpoint: str = "http://localhost:30001"
|
|
) -> dict:
|
|
"""Call the PixelRAG search API."""
|
|
body = json.dumps(
|
|
{"queries": [{"text": query}], "n_docs": min(n_results, 20)}
|
|
).encode()
|
|
req = urllib.request.Request(
|
|
f"{endpoint}/search",
|
|
data=body,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
with urllib.request.urlopen(req, timeout=30) as resp:
|
|
data = json.loads(resp.read())
|
|
|
|
hits = data.get("results", [{}])[0].get("hits", [])
|
|
results = []
|
|
for hit in hits:
|
|
url = hit.get("url", "")
|
|
slug = url.split("/wiki/")[-1] if "/wiki/" in url else ""
|
|
title = slug.replace("_", " ") if slug else url
|
|
results.append(
|
|
{
|
|
"title": title,
|
|
"url": url,
|
|
"score": round(hit["score"], 4),
|
|
"tile": f"tile_{hit.get('tile_index', '?')}_chunk_{hit.get('chunk_index', '?')}",
|
|
}
|
|
)
|
|
return {"query": query, "results": results, "count": len(results)}
|
|
|
|
|
|
def execute_web_fetch(url: str) -> dict:
|
|
"""Fetch text from a URL (simplified — returns first 4000 chars)."""
|
|
req = urllib.request.Request(url, headers={"User-Agent": "PixelRAG-Agent/1.0"})
|
|
with urllib.request.urlopen(req, timeout=15) as resp:
|
|
raw = resp.read().decode("utf-8", errors="replace")
|
|
|
|
# Strip HTML tags for a rough text extraction
|
|
import re
|
|
|
|
text = re.sub(r"<script[^>]*>.*?</script>", "", raw, flags=re.DOTALL)
|
|
text = re.sub(r"<style[^>]*>.*?</style>", "", text, flags=re.DOTALL)
|
|
text = re.sub(r"<[^>]+>", " ", text)
|
|
text = re.sub(r"\s+", " ", text).strip()
|
|
return {"url": url, "content": text[:4000], "truncated": len(text) > 4000}
|
|
|
|
|
|
def handle_tool_call(tool_name: str, tool_input: dict, endpoint: str) -> str:
|
|
"""Execute a tool call and return the result as a string."""
|
|
try:
|
|
if tool_name == "pixelrag_search":
|
|
result = execute_pixelrag_search(
|
|
query=tool_input["query"],
|
|
n_results=tool_input.get("n_results", 5),
|
|
endpoint=endpoint,
|
|
)
|
|
elif tool_name == "web_fetch":
|
|
result = execute_web_fetch(url=tool_input["url"])
|
|
else:
|
|
result = {"error": f"Unknown tool: {tool_name}"}
|
|
except Exception as e:
|
|
result = {"error": str(e)}
|
|
return json.dumps(result)
|
|
|
|
|
|
def run_agent(
|
|
question: str,
|
|
endpoint: str,
|
|
model: str = "claude-sonnet-4-20250514",
|
|
verbose: bool = False,
|
|
) -> str:
|
|
"""Run the agent loop: send question → handle tool calls → return final answer."""
|
|
client = anthropic.Anthropic()
|
|
messages = [{"role": "user", "content": question}]
|
|
|
|
while True:
|
|
response = client.messages.create(
|
|
model=model,
|
|
max_tokens=4096,
|
|
system=SYSTEM_PROMPT,
|
|
tools=TOOLS,
|
|
messages=messages,
|
|
)
|
|
|
|
if verbose:
|
|
print(
|
|
f" [stop_reason={response.stop_reason}, usage={response.usage}]",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
if response.stop_reason == "end_turn":
|
|
# Extract text from response
|
|
text_parts = [b.text for b in response.content if b.type == "text"]
|
|
return "\n".join(text_parts)
|
|
|
|
# Handle tool use
|
|
tool_results = []
|
|
for block in response.content:
|
|
if block.type == "tool_use":
|
|
if verbose:
|
|
print(
|
|
f" [tool: {block.name}({json.dumps(block.input, ensure_ascii=False)})]",
|
|
file=sys.stderr,
|
|
)
|
|
result = handle_tool_call(block.name, block.input, endpoint)
|
|
tool_results.append(
|
|
{
|
|
"type": "tool_result",
|
|
"tool_use_id": block.id,
|
|
"content": result,
|
|
}
|
|
)
|
|
|
|
if not tool_results:
|
|
# No tool calls and not end_turn — shouldn't happen, but handle gracefully
|
|
text_parts = [b.text for b in response.content if b.type == "text"]
|
|
return "\n".join(text_parts) if text_parts else "(no response)"
|
|
|
|
messages.append({"role": "assistant", "content": response.content})
|
|
messages.append({"role": "user", "content": tool_results})
|
|
|
|
|
|
def interactive(endpoint: str, model: str, verbose: bool):
|
|
"""Run interactive conversation loop."""
|
|
print("PixelRAG Agent (Claude + visual search)")
|
|
print(f" endpoint: {endpoint}")
|
|
print(f" model: {model}")
|
|
print(" Type 'quit' to exit.\n")
|
|
|
|
while True:
|
|
try:
|
|
question = input("You: ").strip()
|
|
except (EOFError, KeyboardInterrupt):
|
|
print()
|
|
break
|
|
if not question or question.lower() in ("quit", "exit", "q"):
|
|
break
|
|
|
|
print()
|
|
try:
|
|
answer = run_agent(question, endpoint, model, verbose)
|
|
print(f"Agent: {answer}\n")
|
|
except anthropic.APIError as e:
|
|
print(f"API error: {e}\n")
|
|
except Exception as e:
|
|
print(f"Error: {e}\n")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="PixelRAG Agent — Claude + visual web search"
|
|
)
|
|
parser.add_argument(
|
|
"question", nargs="?", help="Question to ask (omit for interactive mode)"
|
|
)
|
|
parser.add_argument(
|
|
"--endpoint",
|
|
default="http://localhost:30001",
|
|
help="PixelRAG search API endpoint",
|
|
)
|
|
parser.add_argument(
|
|
"--model", default="claude-sonnet-4-20250514", help="Claude model to use"
|
|
)
|
|
parser.add_argument(
|
|
"--verbose", "-v", action="store_true", help="Show tool calls and API details"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if args.question:
|
|
answer = run_agent(args.question, args.endpoint, args.model, args.verbose)
|
|
print(answer)
|
|
else:
|
|
interactive(args.endpoint, args.model, args.verbose)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|