compresr 2.7.2__tar.gz → 2.7.5__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {compresr-2.7.2 → compresr-2.7.5}/PKG-INFO +1 -1
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/engine.py +65 -49
- compresr-2.7.5/compresr/agents/research/__init__.py +24 -0
- compresr-2.7.5/compresr/agents/research/agent.py +334 -0
- compresr-2.7.5/compresr/agents/research/facade.py +94 -0
- compresr-2.7.5/compresr/agents/research/parser.py +62 -0
- compresr-2.7.5/compresr/agents/research/prompts.py +35 -0
- compresr-2.7.5/compresr/agents/research/types.py +53 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/exceptions/__init__.py +1 -6
- {compresr-2.7.2 → compresr-2.7.5}/compresr/exceptions/exceptions.py +0 -14
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/langchain/middleware.py +0 -13
- compresr-2.7.5/compresr/integrations/litellm/__init__.py +54 -0
- compresr-2.7.5/compresr/integrations/litellm/cli.py +41 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/litellm/guardrail.py +233 -181
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/litellm/initializer.py +5 -16
- compresr-2.7.5/compresr/integrations/litellm/shim_installer.py +105 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/litellm/types.py +37 -39
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/llamaindex/memory.py +0 -1
- {compresr-2.7.2 → compresr-2.7.5}/compresr/schemas/usage.py +0 -5
- {compresr-2.7.2 → compresr-2.7.5}/compresr/services/compression.py +12 -8
- {compresr-2.7.2 → compresr-2.7.5}/compresr/services/proxy.py +0 -11
- {compresr-2.7.2 → compresr-2.7.5}/compresr.egg-info/PKG-INFO +1 -1
- {compresr-2.7.2 → compresr-2.7.5}/compresr.egg-info/SOURCES.txt +9 -0
- compresr-2.7.5/compresr.egg-info/entry_points.txt +3 -0
- {compresr-2.7.2 → compresr-2.7.5}/pyproject.toml +13 -18
- compresr-2.7.2/compresr/integrations/litellm/__init__.py +0 -58
- {compresr-2.7.2 → compresr-2.7.5}/LICENSE +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/README.md +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/facades/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/facades/anthropic.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/facades/native.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/facades/openai.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/normalized.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/schemas/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/schemas/anthropic.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/schemas/openai.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/tools/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/agents/tools/web_search.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/clients.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/config.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/_shared/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/_shared/client.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/_shared/compress.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/_shared/errors.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/_shared/filters.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/_shared/kernel.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/_shared/policy.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/_shared/query.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/_shared/tokens.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/langchain/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/langchain/retriever.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/langchain/wrappers.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/langgraph/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/langgraph/checkpoint.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/langgraph/handoff.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/langgraph/nodes.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/langgraph/store.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/llamaindex/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/llamaindex/postprocessor.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/integrations/llamaindex/wrappers.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/schemas/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/schemas/base.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/schemas/inference.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/services/__init__.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr/services/base.py +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr.egg-info/dependency_links.txt +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr.egg-info/requires.txt +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/compresr.egg-info/top_level.txt +0 -0
- {compresr-2.7.2 → compresr-2.7.5}/setup.cfg +0 -0
|
@@ -23,11 +23,6 @@ from .normalized import Citation, CompresrStats, NormalizedResult
|
|
|
23
23
|
logger = logging.getLogger(__name__)
|
|
24
24
|
|
|
25
25
|
|
|
26
|
-
# ---------------------------------------------------------------------------
|
|
27
|
-
# Lazy imports — keep `import compresr.agents` cheap.
|
|
28
|
-
# ---------------------------------------------------------------------------
|
|
29
|
-
|
|
30
|
-
|
|
31
26
|
def _import_lc() -> tuple[Any, Any, Any]:
|
|
32
27
|
"""Resolve ``init_chat_model``, ``create_agent``, and middleware lazily."""
|
|
33
28
|
try:
|
|
@@ -56,28 +51,9 @@ def _CompresrToolMiddleware(*args: Any, **kwargs: Any) -> Any:
|
|
|
56
51
|
return _mw(*args, **kwargs)
|
|
57
52
|
|
|
58
53
|
|
|
59
|
-
# ---------------------------------------------------------------------------
|
|
60
|
-
# Helpers
|
|
61
|
-
# ---------------------------------------------------------------------------
|
|
62
|
-
|
|
63
|
-
|
|
64
54
|
_KNOWN_PROVIDERS = {"anthropic", "openai", "google_genai"}
|
|
65
55
|
|
|
66
56
|
|
|
67
|
-
# ---------------------------------------------------------------------------
|
|
68
|
-
# Solo-WebSearchTool auto-route
|
|
69
|
-
# ---------------------------------------------------------------------------
|
|
70
|
-
#
|
|
71
|
-
# When a generic-agent facade (``client.messages.create`` / ``client.chat
|
|
72
|
-
# .completions.create`` / ``client.run``) is invoked with exactly the
|
|
73
|
-
# Compresr ``WebSearchTool`` and nothing else, the engine silently routes the
|
|
74
|
-
# call through the deep-research loop instead of the standard ``create_agent``
|
|
75
|
-
# middleware path. The caller still gets back their facade's expected shape
|
|
76
|
-
# (Anthropic ``Message`` / OpenAI ``ChatCompletion`` / ``NormalizedResult``),
|
|
77
|
-
# but the loop semantics are research-agent: ``tool_choice="none"`` on the
|
|
78
|
-
# final step, per-snippet compression against the live query, 10-step cap,
|
|
79
|
-
# manual Anthropic ``cache_control`` stamping.
|
|
80
|
-
#
|
|
81
57
|
# Detection is by tool name because ``WebSearchTool.__new__`` returns a
|
|
82
58
|
# ``StructuredTool`` (not a ``WebSearchTool`` instance), so ``isinstance``
|
|
83
59
|
# would always be False.
|
|
@@ -163,6 +139,13 @@ _LLM_BIND_KWARGS = frozenset(
|
|
|
163
139
|
}
|
|
164
140
|
)
|
|
165
141
|
|
|
142
|
+
# Constructor-only kwargs — transport, not per-call LLM knobs. Baked into the
|
|
143
|
+
# chat model at build time so a custom httpx client (corporate proxy / custom
|
|
144
|
+
# CA bundle / verify=False) reaches the provider SDK. Requires langchain-anthropic
|
|
145
|
+
# >= the release that adds ``http_client``/``http_async_client`` to ChatAnthropic;
|
|
146
|
+
# langchain-openai already supports both.
|
|
147
|
+
_LLM_CONSTRUCTOR_KWARGS = frozenset({"http_client", "http_async_client"})
|
|
148
|
+
|
|
166
149
|
|
|
167
150
|
def _parse_llm(llm: str) -> tuple[str, Optional[str]]:
|
|
168
151
|
"""Split ``"provider[:model]"`` (or ``"provider[/model]"``) into
|
|
@@ -259,7 +242,6 @@ def _extract_citations(msg: Any) -> list:
|
|
|
259
242
|
for block in blocks:
|
|
260
243
|
if not isinstance(block, dict):
|
|
261
244
|
continue
|
|
262
|
-
# Anthropic: text block with `citations` list
|
|
263
245
|
for cit in block.get("citations", []) or []:
|
|
264
246
|
if not isinstance(cit, dict):
|
|
265
247
|
continue
|
|
@@ -275,7 +257,6 @@ def _extract_citations(msg: Any) -> list:
|
|
|
275
257
|
provider_metadata=dict(cit),
|
|
276
258
|
)
|
|
277
259
|
)
|
|
278
|
-
# OpenAI Responses: annotations[] with type="url_citation"
|
|
279
260
|
for ann in block.get("annotations", []) or []:
|
|
280
261
|
if not isinstance(ann, dict):
|
|
281
262
|
continue
|
|
@@ -291,7 +272,6 @@ def _extract_citations(msg: Any) -> list:
|
|
|
291
272
|
provider_metadata=dict(ann),
|
|
292
273
|
)
|
|
293
274
|
)
|
|
294
|
-
# Gemini: grounding_metadata on response_metadata
|
|
295
275
|
response_metadata = getattr(msg, "response_metadata", None) or {}
|
|
296
276
|
grounding = response_metadata.get("grounding_metadata") or {}
|
|
297
277
|
for chunk in grounding.get("grounding_chunks", []) or []:
|
|
@@ -446,11 +426,6 @@ def _last_ai_message(messages: list) -> Any:
|
|
|
446
426
|
return messages[-1] if messages else None
|
|
447
427
|
|
|
448
428
|
|
|
449
|
-
# ---------------------------------------------------------------------------
|
|
450
|
-
# _Engine
|
|
451
|
-
# ---------------------------------------------------------------------------
|
|
452
|
-
|
|
453
|
-
|
|
454
429
|
class _Engine:
|
|
455
430
|
"""Private engine shared by the provider-shape facades.
|
|
456
431
|
|
|
@@ -478,6 +453,8 @@ class _Engine:
|
|
|
478
453
|
prompt_cache_ttl: str = "5m",
|
|
479
454
|
prompt_cache_min_messages: int = 2,
|
|
480
455
|
openai_prompt_cache_key: Optional[str] = None,
|
|
456
|
+
llm_http_client: Any = None,
|
|
457
|
+
llm_http_async_client: Any = None,
|
|
481
458
|
) -> None:
|
|
482
459
|
self._compresr_client = compresr_client
|
|
483
460
|
self._provider, self._default_model_name = _parse_llm(llm)
|
|
@@ -487,6 +464,11 @@ class _Engine:
|
|
|
487
464
|
self._prompt_cache_ttl = prompt_cache_ttl
|
|
488
465
|
self._prompt_cache_min_messages = int(prompt_cache_min_messages)
|
|
489
466
|
self._openai_prompt_cache_key = openai_prompt_cache_key
|
|
467
|
+
# Custom httpx clients for the LLM provider transport (corporate proxy /
|
|
468
|
+
# custom CA / verify settings). Baked into every chat-model build; a
|
|
469
|
+
# per-call http_client passed to run() overrides these.
|
|
470
|
+
self._llm_http_client = llm_http_client
|
|
471
|
+
self._llm_http_async_client = llm_http_async_client
|
|
490
472
|
# Chat-model cache keyed by effective model name. Construction is
|
|
491
473
|
# deferred to the first ``run()`` so a provider-only client
|
|
492
474
|
# (``llm="anthropic"``) doesn't try to resolve a model at init time.
|
|
@@ -494,10 +476,6 @@ class _Engine:
|
|
|
494
476
|
# combinations get their own cached chat model.
|
|
495
477
|
self._chat_models: dict[tuple[str, tuple], Any] = {}
|
|
496
478
|
|
|
497
|
-
# ------------------------------------------------------------------
|
|
498
|
-
# Public, read-only accessors
|
|
499
|
-
# ------------------------------------------------------------------
|
|
500
|
-
|
|
501
479
|
@property
|
|
502
480
|
def provider(self) -> str:
|
|
503
481
|
return self._provider
|
|
@@ -516,10 +494,6 @@ class _Engine:
|
|
|
516
494
|
"""
|
|
517
495
|
return self._default_model_name
|
|
518
496
|
|
|
519
|
-
# ------------------------------------------------------------------
|
|
520
|
-
# Build helpers
|
|
521
|
-
# ------------------------------------------------------------------
|
|
522
|
-
|
|
523
497
|
def _get_or_build_chat(self, model_name: str, extra_kwargs: Optional[dict] = None) -> Any:
|
|
524
498
|
"""Return a cached chat model for ``model_name`` + ``extra_kwargs``,
|
|
525
499
|
building on miss.
|
|
@@ -544,6 +518,14 @@ class _Engine:
|
|
|
544
518
|
kwargs: dict = dict(extra_kwargs)
|
|
545
519
|
if self._llm_api_key is not None:
|
|
546
520
|
kwargs.setdefault("api_key", self._llm_api_key)
|
|
521
|
+
# Client-level custom httpx clients. setdefault so a per-call override
|
|
522
|
+
# from run(http_client=...) wins. Only Anthropic and OpenAI chat models
|
|
523
|
+
# accept these kwargs; google_genai uses a different transport.
|
|
524
|
+
if self._provider in ("anthropic", "openai"):
|
|
525
|
+
if self._llm_http_client is not None:
|
|
526
|
+
kwargs.setdefault("http_client", self._llm_http_client)
|
|
527
|
+
if self._llm_http_async_client is not None:
|
|
528
|
+
kwargs.setdefault("http_async_client", self._llm_http_async_client)
|
|
547
529
|
# OpenAI caching is server-side; we only pass ``prompt_cache_key``
|
|
548
530
|
# and ``prompt_cache_retention`` via ``model_kwargs`` so they survive
|
|
549
531
|
# ``bind_tools``.
|
|
@@ -564,8 +546,50 @@ class _Engine:
|
|
|
564
546
|
)
|
|
565
547
|
except TypeError: # pragma: no cover - older langchain-openai
|
|
566
548
|
logger.debug("init_chat_model rejected output_version; falling back without it")
|
|
549
|
+
if self._provider == "anthropic":
|
|
550
|
+
return self._build_anthropic_chat(model_name, kwargs)
|
|
567
551
|
return init_chat_model(f"{self._provider}:{model_name}", **kwargs)
|
|
568
552
|
|
|
553
|
+
def _build_anthropic_chat(self, model_name: str, kwargs: dict) -> Any:
|
|
554
|
+
"""Build an Anthropic chat model, shimming ``http_client`` support on
|
|
555
|
+
versions of langchain-anthropic that predate the native field.
|
|
556
|
+
|
|
557
|
+
Native field present (PR merged / future release) → pass through. Absent
|
|
558
|
+
(current PyPI) → strip the unknown kwargs so the constructor doesn't drop
|
|
559
|
+
them silently, then swap the cached anthropic SDK client for one built
|
|
560
|
+
with the caller's httpx client. The shim is a no-op once native support
|
|
561
|
+
ships, so it can be removed then. See langchain issue #36056."""
|
|
562
|
+
from langchain_anthropic import ChatAnthropic
|
|
563
|
+
|
|
564
|
+
native = "http_client" in ChatAnthropic.model_fields
|
|
565
|
+
sync_client = kwargs.get("http_client")
|
|
566
|
+
async_client = kwargs.get("http_async_client")
|
|
567
|
+
if not native:
|
|
568
|
+
kwargs.pop("http_client", None)
|
|
569
|
+
kwargs.pop("http_async_client", None)
|
|
570
|
+
|
|
571
|
+
chat = init_chat_model(f"anthropic:{model_name}", **kwargs)
|
|
572
|
+
|
|
573
|
+
if not native and (sync_client is not None or async_client is not None):
|
|
574
|
+
self._inject_anthropic_http_client(chat, sync_client, async_client)
|
|
575
|
+
return chat
|
|
576
|
+
|
|
577
|
+
@staticmethod
|
|
578
|
+
def _inject_anthropic_http_client(chat: Any, sync_client: Any, async_client: Any) -> None:
|
|
579
|
+
"""Replace ChatAnthropic's cached SDK clients with ones built around the
|
|
580
|
+
caller's httpx client, reusing langchain's own ``_client_params`` so
|
|
581
|
+
api_key / base_url / retries / headers stay correct. Writes through
|
|
582
|
+
``__dict__`` to seed the ``cached_property`` before first access."""
|
|
583
|
+
import anthropic
|
|
584
|
+
|
|
585
|
+
params = dict(chat._client_params)
|
|
586
|
+
if sync_client is not None:
|
|
587
|
+
chat.__dict__["_client"] = anthropic.Anthropic(**{**params, "http_client": sync_client})
|
|
588
|
+
if async_client is not None:
|
|
589
|
+
chat.__dict__["_async_client"] = anthropic.AsyncClient(
|
|
590
|
+
**{**params, "http_client": async_client}
|
|
591
|
+
)
|
|
592
|
+
|
|
569
593
|
def _openai_cache_model_kwargs(self) -> dict:
|
|
570
594
|
"""Build the OpenAI-specific cache controls. ``prompt_cache_retention``
|
|
571
595
|
is only set for ``ttl="1h"`` so older langchain-openai doesn't reject
|
|
@@ -669,7 +693,7 @@ class _Engine:
|
|
|
669
693
|
"""
|
|
670
694
|
bind_kwargs: dict = {}
|
|
671
695
|
for key in list(kw):
|
|
672
|
-
if key in _LLM_BIND_KWARGS:
|
|
696
|
+
if key in _LLM_BIND_KWARGS or key in _LLM_CONSTRUCTOR_KWARGS:
|
|
673
697
|
bind_kwargs[key] = kw.pop(key)
|
|
674
698
|
if max_tokens is not None:
|
|
675
699
|
# Explicit signature param always reflects customer intent.
|
|
@@ -720,10 +744,6 @@ class _Engine:
|
|
|
720
744
|
# facades layer; for now just keep the original messages.
|
|
721
745
|
return list(messages)
|
|
722
746
|
|
|
723
|
-
# ------------------------------------------------------------------
|
|
724
|
-
# run / arun
|
|
725
|
-
# ------------------------------------------------------------------
|
|
726
|
-
|
|
727
747
|
def run(
|
|
728
748
|
self,
|
|
729
749
|
*,
|
|
@@ -840,10 +860,6 @@ class _Engine:
|
|
|
840
860
|
)
|
|
841
861
|
return self._normalize({"messages": state.messages})
|
|
842
862
|
|
|
843
|
-
# ------------------------------------------------------------------
|
|
844
|
-
# Normalization
|
|
845
|
-
# ------------------------------------------------------------------
|
|
846
|
-
|
|
847
863
|
def _normalize(self, agent_result: Any) -> NormalizedResult:
|
|
848
864
|
"""Reshape ``{"messages": [...]}`` into a :class:`NormalizedResult`.
|
|
849
865
|
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
"""Compresr research agent — multi-step web search with per-snippet compression.
|
|
2
|
+
|
|
3
|
+
Public entry: ``client.research.run(question)`` / ``client.research.search(question)``.
|
|
4
|
+
Loop structure adapted from Perplexity ``search_evals`` (MIT,
|
|
5
|
+
https://github.com/perplexityai/search_evals).
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from .agent import ResearchAgent
|
|
9
|
+
from .facade import ResearchFacade
|
|
10
|
+
from .parser import parse_research_output
|
|
11
|
+
from .prompts import DEFAULT_RESEARCH_SYSTEM_PROMPT
|
|
12
|
+
from .types import Citation, ResearchResult, ResearchUsage, Step, StepKind
|
|
13
|
+
|
|
14
|
+
__all__ = [
|
|
15
|
+
"Citation",
|
|
16
|
+
"DEFAULT_RESEARCH_SYSTEM_PROMPT",
|
|
17
|
+
"ResearchAgent",
|
|
18
|
+
"ResearchFacade",
|
|
19
|
+
"ResearchResult",
|
|
20
|
+
"ResearchUsage",
|
|
21
|
+
"Step",
|
|
22
|
+
"StepKind",
|
|
23
|
+
"parse_research_output",
|
|
24
|
+
]
|
|
@@ -0,0 +1,334 @@
|
|
|
1
|
+
"""``ResearchAgent`` — strict-output ReAct loop with per-snippet compression.
|
|
2
|
+
|
|
3
|
+
Bypasses ``_Engine.run`` so we can control ``tool_choice`` per step (forces
|
|
4
|
+
the model to commit on the final step) and compress each tool result against
|
|
5
|
+
the live query before it enters the conversation. Cache control is stamped
|
|
6
|
+
manually for Anthropic since ``create_agent`` is not in the call path.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import logging
|
|
12
|
+
import re
|
|
13
|
+
import time
|
|
14
|
+
from dataclasses import dataclass, field
|
|
15
|
+
from typing import Any, Optional, Sequence
|
|
16
|
+
|
|
17
|
+
from .parser import parse_research_output
|
|
18
|
+
from .prompts import DEFAULT_RESEARCH_SYSTEM_PROMPT
|
|
19
|
+
from .types import Citation, ResearchResult, ResearchUsage, Step
|
|
20
|
+
|
|
21
|
+
logger = logging.getLogger(__name__)
|
|
22
|
+
|
|
23
|
+
_DEFAULT_MAX_CONTEXT_TOKENS = 120_000
|
|
24
|
+
_DEFAULT_MAX_STEPS = 10
|
|
25
|
+
_DEFAULT_MIN_COMPRESS_TOKENS = 100
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def _estimate_tokens(text: str) -> int:
|
|
29
|
+
return max(1, len(text) // 4)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
@dataclass
|
|
33
|
+
class _LoopState:
|
|
34
|
+
"""Internal result of :meth:`ResearchAgent._run_loop`.
|
|
35
|
+
|
|
36
|
+
Exposed to the engine's solo-WebSearchTool fast-path so it can feed the
|
|
37
|
+
raw message chain to :meth:`_Engine._normalize` while the public
|
|
38
|
+
:meth:`run` keeps producing a typed :class:`ResearchResult`.
|
|
39
|
+
"""
|
|
40
|
+
|
|
41
|
+
messages: list
|
|
42
|
+
ai_history: list = field(default_factory=list)
|
|
43
|
+
trajectory: list = field(default_factory=list)
|
|
44
|
+
last_ai: Optional[Any] = None
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
class ResearchAgent:
|
|
48
|
+
"""Multi-step web-research agent with on-the-fly compression.
|
|
49
|
+
|
|
50
|
+
Args:
|
|
51
|
+
engine: A live :class:`compresr.agents.engine._Engine`.
|
|
52
|
+
search_tool: LangChain ``BaseTool`` taking ``query: str``.
|
|
53
|
+
max_steps: Hard cap on LLM turns. On the final turn ``tool_choice``
|
|
54
|
+
is set to ``"none"`` so the model has to commit.
|
|
55
|
+
system_prompt: Override the default research prompt.
|
|
56
|
+
compress_snippets: Compress each tool result via
|
|
57
|
+
``client.compress(query=live_query)`` before appending.
|
|
58
|
+
compression_model: Compresr model name (default ``"latte_v1"``).
|
|
59
|
+
min_compress_tokens: Skip compression for tool results below this size.
|
|
60
|
+
max_context_tokens: Drop the oldest tool-result pair when context
|
|
61
|
+
exceeds this cap.
|
|
62
|
+
"""
|
|
63
|
+
|
|
64
|
+
def __init__(
|
|
65
|
+
self,
|
|
66
|
+
*,
|
|
67
|
+
engine: Any,
|
|
68
|
+
search_tool: Any,
|
|
69
|
+
max_steps: int = _DEFAULT_MAX_STEPS,
|
|
70
|
+
system_prompt: Optional[str] = None,
|
|
71
|
+
compress_snippets: bool = True,
|
|
72
|
+
compression_model: str = "latte_v1",
|
|
73
|
+
min_compress_tokens: int = _DEFAULT_MIN_COMPRESS_TOKENS,
|
|
74
|
+
max_context_tokens: int = _DEFAULT_MAX_CONTEXT_TOKENS,
|
|
75
|
+
) -> None:
|
|
76
|
+
self._engine = engine
|
|
77
|
+
self._search_tool = search_tool
|
|
78
|
+
self._max_steps = max(1, int(max_steps))
|
|
79
|
+
self._system_prompt = system_prompt or DEFAULT_RESEARCH_SYSTEM_PROMPT
|
|
80
|
+
self._compress_snippets = bool(compress_snippets)
|
|
81
|
+
self._compression_model = compression_model
|
|
82
|
+
self._min_compress_tokens = int(min_compress_tokens)
|
|
83
|
+
self._max_context_tokens = int(max_context_tokens)
|
|
84
|
+
|
|
85
|
+
def run(self, question: str, *, model: Optional[str] = None) -> ResearchResult:
|
|
86
|
+
state = self._run_loop(question, model=model)
|
|
87
|
+
usage = self._aggregate_usage(
|
|
88
|
+
state.ai_history,
|
|
89
|
+
search_calls=sum(1 for s in state.trajectory if s.type == "search"),
|
|
90
|
+
)
|
|
91
|
+
final_text = self._stringify(state.last_ai.content) if state.last_ai is not None else ""
|
|
92
|
+
parsed = parse_research_output(final_text)
|
|
93
|
+
citations = self._collect_citations(parsed.citation_urls, state.trajectory)
|
|
94
|
+
|
|
95
|
+
return ResearchResult(
|
|
96
|
+
answer=parsed.answer,
|
|
97
|
+
explanation=parsed.explanation,
|
|
98
|
+
confidence=parsed.confidence,
|
|
99
|
+
text=final_text,
|
|
100
|
+
citations=citations,
|
|
101
|
+
trajectory=state.trajectory,
|
|
102
|
+
usage=usage,
|
|
103
|
+
raw=state.last_ai,
|
|
104
|
+
)
|
|
105
|
+
|
|
106
|
+
def _run_loop(
|
|
107
|
+
self,
|
|
108
|
+
question: str,
|
|
109
|
+
*,
|
|
110
|
+
model: Optional[str] = None,
|
|
111
|
+
extra_chat_kwargs: Optional[dict] = None,
|
|
112
|
+
) -> _LoopState:
|
|
113
|
+
"""Execute the strict-output ReAct loop and return raw state.
|
|
114
|
+
|
|
115
|
+
Shared by :meth:`run` (which parses ``ResearchResult``) and the
|
|
116
|
+
engine's solo-WebSearchTool fast-path (which feeds ``state.messages``
|
|
117
|
+
into :meth:`_Engine._normalize`).
|
|
118
|
+
|
|
119
|
+
``extra_chat_kwargs`` lets the engine path bake per-call LLM knobs
|
|
120
|
+
(``max_tokens``, ``temperature``…) into the cached chat model so
|
|
121
|
+
they survive ``bind_tools``.
|
|
122
|
+
"""
|
|
123
|
+
from langchain_core.messages import ( # type: ignore[import-not-found]
|
|
124
|
+
HumanMessage,
|
|
125
|
+
SystemMessage,
|
|
126
|
+
ToolMessage,
|
|
127
|
+
)
|
|
128
|
+
|
|
129
|
+
effective_model = self._engine._resolve_model(model)
|
|
130
|
+
chat = self._engine._get_or_build_chat(effective_model, extra_kwargs=extra_chat_kwargs)
|
|
131
|
+
|
|
132
|
+
messages: list = [
|
|
133
|
+
SystemMessage(content=self._system_prompt),
|
|
134
|
+
HumanMessage(content=question),
|
|
135
|
+
]
|
|
136
|
+
ai_history: list = []
|
|
137
|
+
trajectory: list[Step] = []
|
|
138
|
+
last_ai: Optional[Any] = None
|
|
139
|
+
|
|
140
|
+
for step_idx in range(self._max_steps):
|
|
141
|
+
self._maybe_truncate(messages)
|
|
142
|
+
|
|
143
|
+
tool_choice = "none" if step_idx == self._max_steps - 1 else "auto"
|
|
144
|
+
bound = chat.bind_tools([self._search_tool], tool_choice=tool_choice)
|
|
145
|
+
invoke_messages = self._apply_cache_control(messages)
|
|
146
|
+
|
|
147
|
+
t0 = time.time()
|
|
148
|
+
try:
|
|
149
|
+
response = bound.invoke(invoke_messages)
|
|
150
|
+
except Exception as exc: # noqa: BLE001
|
|
151
|
+
trajectory.append(
|
|
152
|
+
Step(type="error", text=str(exc), latency_s=round(time.time() - t0, 3))
|
|
153
|
+
)
|
|
154
|
+
break
|
|
155
|
+
step_latency = round(time.time() - t0, 3)
|
|
156
|
+
|
|
157
|
+
ai_history.append(response)
|
|
158
|
+
last_ai = response
|
|
159
|
+
messages.append(response)
|
|
160
|
+
|
|
161
|
+
tool_calls = getattr(response, "tool_calls", None) or []
|
|
162
|
+
if not tool_calls:
|
|
163
|
+
trajectory.append(
|
|
164
|
+
Step(
|
|
165
|
+
type="answer",
|
|
166
|
+
text=self._stringify(response.content),
|
|
167
|
+
latency_s=step_latency,
|
|
168
|
+
)
|
|
169
|
+
)
|
|
170
|
+
break
|
|
171
|
+
|
|
172
|
+
for tc in tool_calls:
|
|
173
|
+
args = tc.get("args") if isinstance(tc, dict) else {}
|
|
174
|
+
args = args if isinstance(args, dict) else {}
|
|
175
|
+
query = args.get("query")
|
|
176
|
+
trajectory.append(Step(type="search", query=query, latency_s=step_latency))
|
|
177
|
+
|
|
178
|
+
raw_result = self._invoke_search(args)
|
|
179
|
+
final_result = (
|
|
180
|
+
self._compress(raw_result, query or question)
|
|
181
|
+
if self._compress_snippets
|
|
182
|
+
else raw_result
|
|
183
|
+
)
|
|
184
|
+
trajectory.append(
|
|
185
|
+
Step(
|
|
186
|
+
type="tool_result",
|
|
187
|
+
chars=len(final_result),
|
|
188
|
+
text=final_result,
|
|
189
|
+
)
|
|
190
|
+
)
|
|
191
|
+
|
|
192
|
+
tool_call_id = tc.get("id") if isinstance(tc, dict) else None
|
|
193
|
+
messages.append(
|
|
194
|
+
ToolMessage(
|
|
195
|
+
content=final_result,
|
|
196
|
+
tool_call_id=tool_call_id or "",
|
|
197
|
+
name=getattr(self._search_tool, "name", "search_web"),
|
|
198
|
+
)
|
|
199
|
+
)
|
|
200
|
+
|
|
201
|
+
return _LoopState(
|
|
202
|
+
messages=messages,
|
|
203
|
+
ai_history=ai_history,
|
|
204
|
+
trajectory=trajectory,
|
|
205
|
+
last_ai=last_ai,
|
|
206
|
+
)
|
|
207
|
+
|
|
208
|
+
def _apply_cache_control(self, messages: list) -> list:
|
|
209
|
+
"""Stamp Anthropic ephemeral cache_control on the last message.
|
|
210
|
+
|
|
211
|
+
Mirrors ``AnthropicPromptCachingMiddleware`` for the bypass path —
|
|
212
|
+
the research agent does not go through ``create_agent`` so the
|
|
213
|
+
engine's middleware never runs. No-op for non-Anthropic providers.
|
|
214
|
+
"""
|
|
215
|
+
provider = getattr(self._engine, "_provider", None)
|
|
216
|
+
if provider != "anthropic":
|
|
217
|
+
return messages
|
|
218
|
+
if not getattr(self._engine, "_enable_prompt_cache", False):
|
|
219
|
+
return messages
|
|
220
|
+
min_msgs = getattr(self._engine, "_prompt_cache_min_messages", 2)
|
|
221
|
+
if len(messages) < min_msgs:
|
|
222
|
+
return messages
|
|
223
|
+
|
|
224
|
+
ttl = getattr(self._engine, "_prompt_cache_ttl", "5m")
|
|
225
|
+
marker = {"type": "ephemeral", "ttl": ttl}
|
|
226
|
+
last = messages[-1]
|
|
227
|
+
content = getattr(last, "content", None)
|
|
228
|
+
if isinstance(content, str):
|
|
229
|
+
new_content: list = [{"type": "text", "text": content, "cache_control": marker}]
|
|
230
|
+
elif isinstance(content, list) and content:
|
|
231
|
+
new_content = list(content)
|
|
232
|
+
tail = new_content[-1]
|
|
233
|
+
if isinstance(tail, dict):
|
|
234
|
+
tail = dict(tail)
|
|
235
|
+
tail["cache_control"] = marker
|
|
236
|
+
new_content[-1] = tail
|
|
237
|
+
else:
|
|
238
|
+
new_content.append({"type": "text", "text": str(tail), "cache_control": marker})
|
|
239
|
+
else:
|
|
240
|
+
return messages
|
|
241
|
+
|
|
242
|
+
try:
|
|
243
|
+
patched = last.model_copy(update={"content": new_content})
|
|
244
|
+
except Exception: # noqa: BLE001
|
|
245
|
+
patched = type(last)(content=new_content)
|
|
246
|
+
return [*messages[:-1], patched]
|
|
247
|
+
|
|
248
|
+
def _invoke_search(self, args: dict) -> str:
|
|
249
|
+
try:
|
|
250
|
+
out = self._search_tool.invoke(args)
|
|
251
|
+
except Exception as exc: # noqa: BLE001
|
|
252
|
+
logger.warning("research: search tool raised %s; agent will continue", exc)
|
|
253
|
+
return f"Search error: {exc}"
|
|
254
|
+
return self._stringify(out)
|
|
255
|
+
|
|
256
|
+
def _compress(self, text: str, query: str) -> str:
|
|
257
|
+
if not text or _estimate_tokens(text) < self._min_compress_tokens:
|
|
258
|
+
return text
|
|
259
|
+
client = self._engine._compresr_client
|
|
260
|
+
try:
|
|
261
|
+
resp = client.compress(
|
|
262
|
+
context=text,
|
|
263
|
+
query=query,
|
|
264
|
+
compression_model_name=self._compression_model,
|
|
265
|
+
)
|
|
266
|
+
data = getattr(resp, "data", None)
|
|
267
|
+
compressed = getattr(data, "compressed_context", None) if data is not None else None
|
|
268
|
+
return compressed if isinstance(compressed, str) and compressed else text
|
|
269
|
+
except Exception as exc: # noqa: BLE001
|
|
270
|
+
logger.debug("research: compression failed (%s); using raw snippet", exc)
|
|
271
|
+
return text
|
|
272
|
+
|
|
273
|
+
def _maybe_truncate(self, messages: list) -> None:
|
|
274
|
+
if not self._max_context_tokens:
|
|
275
|
+
return
|
|
276
|
+
total = sum(_estimate_tokens(self._stringify(getattr(m, "content", ""))) for m in messages)
|
|
277
|
+
if total <= self._max_context_tokens:
|
|
278
|
+
return
|
|
279
|
+
guard = 0
|
|
280
|
+
while total > self._max_context_tokens and len(messages) >= 4 and guard < 4:
|
|
281
|
+
del messages[2:4]
|
|
282
|
+
total = sum(
|
|
283
|
+
_estimate_tokens(self._stringify(getattr(m, "content", ""))) for m in messages
|
|
284
|
+
)
|
|
285
|
+
guard += 1
|
|
286
|
+
|
|
287
|
+
def _aggregate_usage(
|
|
288
|
+
self,
|
|
289
|
+
ai_messages: Sequence[Any],
|
|
290
|
+
*,
|
|
291
|
+
search_calls: int = 0,
|
|
292
|
+
) -> ResearchUsage:
|
|
293
|
+
from compresr.agents.engine import _aggregate_usage
|
|
294
|
+
|
|
295
|
+
agg = _aggregate_usage(list(ai_messages))
|
|
296
|
+
return ResearchUsage(
|
|
297
|
+
input_tokens=int(agg.get("input_tokens", 0) or 0),
|
|
298
|
+
output_tokens=int(agg.get("output_tokens", 0) or 0),
|
|
299
|
+
cache_read_tokens=int(agg.get("cache_read_input_tokens", 0) or 0),
|
|
300
|
+
cache_creation_tokens=int(agg.get("cache_creation_input_tokens", 0) or 0),
|
|
301
|
+
calls=len(ai_messages),
|
|
302
|
+
search_calls=int(search_calls or 0),
|
|
303
|
+
)
|
|
304
|
+
|
|
305
|
+
def _collect_citations(
|
|
306
|
+
self, parsed_urls: Sequence[str], trajectory: Sequence[Step]
|
|
307
|
+
) -> list[Citation]:
|
|
308
|
+
urls: list[str] = []
|
|
309
|
+
for u in parsed_urls:
|
|
310
|
+
if u not in urls:
|
|
311
|
+
urls.append(u)
|
|
312
|
+
for step in trajectory:
|
|
313
|
+
if step.type != "tool_result" or not step.text:
|
|
314
|
+
continue
|
|
315
|
+
for url in re.findall(r"https?://[^\s,;<>\"')]+", step.text):
|
|
316
|
+
if url not in urls:
|
|
317
|
+
urls.append(url)
|
|
318
|
+
return [Citation(url=u) for u in urls]
|
|
319
|
+
|
|
320
|
+
@staticmethod
|
|
321
|
+
def _stringify(content: Any) -> str:
|
|
322
|
+
if isinstance(content, str):
|
|
323
|
+
return content
|
|
324
|
+
if isinstance(content, list):
|
|
325
|
+
parts = []
|
|
326
|
+
for block in content:
|
|
327
|
+
if isinstance(block, dict):
|
|
328
|
+
text = block.get("text")
|
|
329
|
+
if isinstance(text, str):
|
|
330
|
+
parts.append(text)
|
|
331
|
+
elif isinstance(block, str):
|
|
332
|
+
parts.append(block)
|
|
333
|
+
return "".join(parts)
|
|
334
|
+
return str(content) if content is not None else ""
|
|
@@ -0,0 +1,94 @@
|
|
|
1
|
+
"""``client.research`` facade — thin wrapper over :class:`ResearchAgent`."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import os
|
|
6
|
+
from typing import Any, Optional
|
|
7
|
+
|
|
8
|
+
from .agent import ResearchAgent
|
|
9
|
+
from .types import ResearchResult
|
|
10
|
+
|
|
11
|
+
_SEARCH_PROVIDER_ENV = {
|
|
12
|
+
"tavily": ("TAVILY_API_KEY",),
|
|
13
|
+
"brave": ("BRAVE_SEARCH_API_KEY", "BRAVE_API_KEY"),
|
|
14
|
+
}
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class ResearchFacade:
|
|
18
|
+
def __init__(self, engine: Any) -> None:
|
|
19
|
+
self._engine = engine
|
|
20
|
+
|
|
21
|
+
def run(
|
|
22
|
+
self,
|
|
23
|
+
question: str,
|
|
24
|
+
*,
|
|
25
|
+
search: Any = "tavily",
|
|
26
|
+
max_steps: int = 10,
|
|
27
|
+
model: Optional[str] = None,
|
|
28
|
+
compress_snippets: bool = True,
|
|
29
|
+
compression_model: str = "latte_v1",
|
|
30
|
+
min_compress_tokens: int = 100,
|
|
31
|
+
max_context_tokens: int = 120_000,
|
|
32
|
+
system_prompt: Optional[str] = None,
|
|
33
|
+
) -> ResearchResult:
|
|
34
|
+
"""Multi-step research loop.
|
|
35
|
+
|
|
36
|
+
``search`` is either a provider name (``"tavily"`` | ``"brave"``)
|
|
37
|
+
or a ready-made LangChain ``BaseTool``.
|
|
38
|
+
"""
|
|
39
|
+
tool = self._resolve_search_tool(search)
|
|
40
|
+
agent = ResearchAgent(
|
|
41
|
+
engine=self._engine,
|
|
42
|
+
search_tool=tool,
|
|
43
|
+
max_steps=max_steps,
|
|
44
|
+
system_prompt=system_prompt,
|
|
45
|
+
compress_snippets=compress_snippets,
|
|
46
|
+
compression_model=compression_model,
|
|
47
|
+
min_compress_tokens=min_compress_tokens,
|
|
48
|
+
max_context_tokens=max_context_tokens,
|
|
49
|
+
)
|
|
50
|
+
return agent.run(question, model=model)
|
|
51
|
+
|
|
52
|
+
def search(
|
|
53
|
+
self,
|
|
54
|
+
question: str,
|
|
55
|
+
*,
|
|
56
|
+
search: Any = "tavily",
|
|
57
|
+
model: Optional[str] = None,
|
|
58
|
+
compress_snippets: bool = True,
|
|
59
|
+
compression_model: str = "latte_v1",
|
|
60
|
+
min_compress_tokens: int = 100,
|
|
61
|
+
max_context_tokens: int = 120_000,
|
|
62
|
+
system_prompt: Optional[str] = None,
|
|
63
|
+
) -> ResearchResult:
|
|
64
|
+
"""Single-shot: one search + forced answer (``max_steps=2``)."""
|
|
65
|
+
return self.run(
|
|
66
|
+
question,
|
|
67
|
+
search=search,
|
|
68
|
+
max_steps=2,
|
|
69
|
+
model=model,
|
|
70
|
+
compress_snippets=compress_snippets,
|
|
71
|
+
compression_model=compression_model,
|
|
72
|
+
min_compress_tokens=min_compress_tokens,
|
|
73
|
+
max_context_tokens=max_context_tokens,
|
|
74
|
+
system_prompt=system_prompt,
|
|
75
|
+
)
|
|
76
|
+
|
|
77
|
+
def _resolve_search_tool(self, search: Any) -> Any:
|
|
78
|
+
if not isinstance(search, str):
|
|
79
|
+
return search
|
|
80
|
+
provider = search.lower()
|
|
81
|
+
if provider not in _SEARCH_PROVIDER_ENV:
|
|
82
|
+
raise ValueError(
|
|
83
|
+
f"unsupported search provider {search!r}; expected one of "
|
|
84
|
+
f"{sorted(_SEARCH_PROVIDER_ENV)} or a LangChain BaseTool"
|
|
85
|
+
)
|
|
86
|
+
from ..tools.web_search import WebSearchTool
|
|
87
|
+
|
|
88
|
+
api_key: Optional[str] = None
|
|
89
|
+
for env_var in _SEARCH_PROVIDER_ENV[provider]:
|
|
90
|
+
v = os.environ.get(env_var)
|
|
91
|
+
if v:
|
|
92
|
+
api_key = v
|
|
93
|
+
break
|
|
94
|
+
return WebSearchTool(provider=provider, api_key=api_key)
|