192 lines
7.0 KiB
Python
192 lines
7.0 KiB
Python
from __future__ import annotations
|
||
import logging
|
||
import time
|
||
|
||
import httpx
|
||
|
||
from app.models.agent import AgentPersona
|
||
from app.models.world import WorldState
|
||
from app.services.llm import LLMClient
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
RESEARCH_SUMMARIZE_SYSTEM = """You are {name}, a {age}-year-old {role}.
|
||
Personality: honesty={honesty:.1f}, ambition={ambition:.1f}, empathy={empathy:.1f}, conformity={conformity:.1f}
|
||
{market_personality}
|
||
|
||
Your current beliefs:
|
||
{beliefs}
|
||
|
||
You just searched the internet for: "{query}"
|
||
Reason: {reason}
|
||
|
||
Here are the search results:
|
||
{search_results}
|
||
|
||
Based on YOUR personality and beliefs, what do you take away from this?
|
||
Write 2-3 sentences of what you learned and how it affects your thinking.
|
||
Be specific — cite numbers, names, or facts you found.
|
||
If the results contradict your beliefs, note the tension.
|
||
If they confirm your beliefs, note the reinforcement."""
|
||
|
||
|
||
class ResearchService:
|
||
def __init__(
|
||
self,
|
||
llm: LLMClient,
|
||
enabled: bool = True,
|
||
max_per_round: int = 5,
|
||
searxng_base_url: str = "",
|
||
searxng_token: str = "",
|
||
searxng_timeout: float = 15.0,
|
||
searxng_language: str = "auto",
|
||
searxng_safesearch: int = 1,
|
||
cache_ttl_seconds: int = 300,
|
||
):
|
||
self.llm = llm
|
||
self.enabled = enabled
|
||
self.max_per_round = max_per_round
|
||
self.searxng_base_url = searxng_base_url.rstrip("/")
|
||
self.searxng_token = searxng_token
|
||
self.searxng_timeout = max(1.0, float(searxng_timeout))
|
||
self.searxng_language = searxng_language or "auto"
|
||
self.searxng_safesearch = max(0, min(2, int(searxng_safesearch)))
|
||
self.cache_ttl_seconds = max(0, int(cache_ttl_seconds))
|
||
self._client: httpx.AsyncClient | None = None
|
||
self._cache: dict[str, tuple[float, list[dict]]] = {}
|
||
|
||
async def _get_client(self) -> httpx.AsyncClient:
|
||
if self._client is None or self._client.is_closed:
|
||
timeout = httpx.Timeout(self.searxng_timeout, connect=min(5.0, self.searxng_timeout))
|
||
self._client = httpx.AsyncClient(timeout=timeout, follow_redirects=True)
|
||
return self._client
|
||
|
||
async def aclose(self) -> None:
|
||
if self._client is not None:
|
||
await self._client.aclose()
|
||
self._client = None
|
||
|
||
def _cache_key(self, query: str, max_results: int) -> str:
|
||
return "|".join((query.strip().casefold(), str(max_results), self.searxng_language, str(self.searxng_safesearch)))
|
||
|
||
async def _raw_search(self, query: str, max_results: int = 5) -> list[dict]:
|
||
"""Search SearXNG and normalize its response to the legacy shape.
|
||
|
||
Search failures intentionally return an empty list. Research is an
|
||
optional action and must not abort an otherwise healthy simulation.
|
||
"""
|
||
query = (query or "").strip()
|
||
if not query:
|
||
return []
|
||
if not self.searxng_base_url:
|
||
logger.warning("SearXNG 未配置，跳过搜索（query=%s）", query[:120])
|
||
return []
|
||
|
||
max_results = max(1, min(20, int(max_results)))
|
||
cache_key = self._cache_key(query, max_results)
|
||
now = time.monotonic()
|
||
cached = self._cache.get(cache_key)
|
||
if self.cache_ttl_seconds > 0 and cached and now - cached[0] < self.cache_ttl_seconds:
|
||
return [dict(item) for item in cached[1]]
|
||
if cached and self.cache_ttl_seconds > 0:
|
||
self._cache.pop(cache_key, None)
|
||
|
||
try:
|
||
client = await self._get_client()
|
||
params = {
|
||
"q": query,
|
||
"format": "json",
|
||
"language": self.searxng_language,
|
||
"safesearch": self.searxng_safesearch,
|
||
}
|
||
if self.searxng_token:
|
||
params["token"] = self.searxng_token
|
||
response = await client.get(f"{self.searxng_base_url}/search", params=params)
|
||
response.raise_for_status()
|
||
payload = response.json()
|
||
raw_results = payload.get("results", []) if isinstance(payload, dict) else []
|
||
results = []
|
||
for result in raw_results[:max_results]:
|
||
if not isinstance(result, dict):
|
||
continue
|
||
results.append({
|
||
"title": str(result.get("title") or ""),
|
||
"body": str(result.get("content") or ""),
|
||
"href": str(result.get("url") or ""),
|
||
})
|
||
if self.cache_ttl_seconds > 0:
|
||
self._cache[cache_key] = (now, results)
|
||
if len(self._cache) > 256:
|
||
oldest = min(self._cache, key=lambda key: self._cache[key][0])
|
||
self._cache.pop(oldest, None)
|
||
return results
|
||
except (httpx.HTTPError, ValueError) as e:
|
||
logger.warning("SearXNG search failed for '%s': %s", query[:120], e)
|
||
return []
|
||
|
||
async def search_and_summarize(
|
||
self,
|
||
query: str,
|
||
reason: str,
|
||
agent: AgentPersona,
|
||
world_state: WorldState,
|
||
) -> str:
|
||
if not self.enabled:
|
||
return ""
|
||
|
||
results = await self._raw_search(query)
|
||
if not results:
|
||
return ""
|
||
|
||
snippets = []
|
||
for i, r in enumerate(results[:5], 1):
|
||
title = r.get("title", "")
|
||
body = r.get("body", "")
|
||
href = r.get("href", "")
|
||
snippets.append(f"{i}. {title}\n {body}\n Source: {href}")
|
||
|
||
search_results_text = "\n\n".join(snippets)
|
||
|
||
p = agent.personality
|
||
market_personality = ""
|
||
has_market_traits = (
|
||
p.brand_loyalty != 0.5 or p.price_sensitivity != 0.5
|
||
or p.social_proof != 0.5 or p.novelty_seeking != 0.5
|
||
)
|
||
if has_market_traits:
|
||
market_personality = (
|
||
f"Consumer traits: brand_loyalty={p.brand_loyalty:.1f}, "
|
||
f"price_sensitivity={p.price_sensitivity:.1f}, "
|
||
f"social_proof={p.social_proof:.1f}, "
|
||
f"novelty_seeking={p.novelty_seeking:.1f}"
|
||
)
|
||
|
||
beliefs = "\n".join(f"- {b}" for b in agent.beliefs[:5]) or "Still forming opinions."
|
||
|
||
system = RESEARCH_SUMMARIZE_SYSTEM.format(
|
||
name=agent.name,
|
||
age=agent.age,
|
||
role=agent.role,
|
||
honesty=p.honesty,
|
||
ambition=p.ambition,
|
||
empathy=p.empathy,
|
||
conformity=p.conformity,
|
||
market_personality=market_personality,
|
||
beliefs=beliefs,
|
||
query=query,
|
||
reason=reason,
|
||
search_results=search_results_text,
|
||
)
|
||
|
||
try:
|
||
digest = await self.llm.generate(
|
||
system=system,
|
||
user="What did you learn? Summarize in 2-3 sentences.",
|
||
json_mode=False,
|
||
max_tokens=200,
|
||
)
|
||
return digest.strip()
|
||
except Exception as e:
|
||
logger.warning("Research summarization failed for agent %s: %s", agent.name, e)
|
||
return ""
|