feat(integrations): add AG2 framework integration (#720)
Add hindsight-ag2 package providing persistent memory tools for AG2 agents via retain/recall/reflect operations.
This commit is contained in:
parent
62c0992075
commit
731238707d
10 changed files with 2759 additions and 0 deletions
143
hindsight-integrations/ag2/README.md
Normal file
143
hindsight-integrations/ag2/README.md
Normal file
|
|
@ -0,0 +1,143 @@
|
||||||
|
# hindsight-ag2
|
||||||
|
|
||||||
|
AG2 integration for [Hindsight](https://github.com/vectorize-io/hindsight) — persistent long-term memory for AI agents.
|
||||||
|
|
||||||
|
Provides Hindsight-backed tool functions that give [AG2](https://ag2.ai) agents long-term memory across conversations via retain/recall/reflect operations.
|
||||||
|
|
||||||
|
## Prerequisites
|
||||||
|
|
||||||
|
- Python 3.10+
|
||||||
|
- Running Hindsight instance ([quickstart](https://github.com/vectorize-io/hindsight#quick-start))
|
||||||
|
|
||||||
|
## Installation
|
||||||
|
|
||||||
|
```bash
|
||||||
|
pip install hindsight-ag2
|
||||||
|
```
|
||||||
|
|
||||||
|
## Quick Start
|
||||||
|
|
||||||
|
```python
|
||||||
|
from autogen import AssistantAgent, UserProxyAgent, LLMConfig
|
||||||
|
from hindsight_ag2 import register_hindsight_tools
|
||||||
|
|
||||||
|
llm_config = LLMConfig(api_type="openai", model="gpt-4o-mini")
|
||||||
|
|
||||||
|
with llm_config:
|
||||||
|
assistant = AssistantAgent(
|
||||||
|
name="assistant",
|
||||||
|
system_message="You are a helpful assistant with long-term memory.",
|
||||||
|
)
|
||||||
|
user_proxy = UserProxyAgent(
|
||||||
|
name="user",
|
||||||
|
human_input_mode="NEVER",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Register Hindsight memory tools on both agents
|
||||||
|
register_hindsight_tools(
|
||||||
|
assistant, user_proxy,
|
||||||
|
bank_id="my-bank",
|
||||||
|
hindsight_api_url="http://localhost:8888",
|
||||||
|
)
|
||||||
|
|
||||||
|
# The assistant can now use hindsight_retain, hindsight_recall, hindsight_reflect
|
||||||
|
result = user_proxy.initiate_chat(
|
||||||
|
assistant,
|
||||||
|
message="Remember that I prefer Python over JavaScript.",
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
## Tools
|
||||||
|
|
||||||
|
| Tool | Operation | Description |
|
||||||
|
|------|-----------|-------------|
|
||||||
|
| `hindsight_retain` | Retain | Store facts, preferences, decisions to long-term memory |
|
||||||
|
| `hindsight_recall` | Recall | Multi-strategy search across stored memories |
|
||||||
|
| `hindsight_reflect` | Reflect | Synthesize reasoned answers from memories |
|
||||||
|
|
||||||
|
## Configuration
|
||||||
|
|
||||||
|
### Global config
|
||||||
|
|
||||||
|
```python
|
||||||
|
from hindsight_ag2 import configure
|
||||||
|
|
||||||
|
configure(
|
||||||
|
hindsight_api_url="http://localhost:8888",
|
||||||
|
api_key="your-key", # or set HINDSIGHT_API_KEY env var
|
||||||
|
budget="mid", # low / mid / high
|
||||||
|
max_tokens=4096,
|
||||||
|
tags=["source:ag2"], # default tags for retain
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
### Per-call overrides
|
||||||
|
|
||||||
|
All global settings can be overridden per `create_hindsight_tools()` call:
|
||||||
|
|
||||||
|
| Parameter | Description | Default |
|
||||||
|
|-----------|-------------|---------|
|
||||||
|
| `bank_id` | Memory bank ID (required) | — |
|
||||||
|
| `client` | Pre-configured `Hindsight` client | — |
|
||||||
|
| `hindsight_api_url` | API URL | Global config or production |
|
||||||
|
| `api_key` | API key | Global config or env var |
|
||||||
|
| `budget` | Recall/reflect budget | `"mid"` |
|
||||||
|
| `max_tokens` | Max tokens for recall | `4096` |
|
||||||
|
| `tags` | Tags for retain operations | `None` |
|
||||||
|
| `recall_tags` | Tags to filter recall | `None` |
|
||||||
|
| `recall_tags_match` | Tag match mode | `"any"` |
|
||||||
|
| `retain_metadata` | Metadata dict for retain | `None` |
|
||||||
|
| `retain_document_id` | Document ID for retain | `None` |
|
||||||
|
| `recall_types` | Fact types to filter | `None` |
|
||||||
|
| `recall_include_entities` | Include entities in recall | `False` |
|
||||||
|
| `reflect_context` | Additional context for reflect | `None` |
|
||||||
|
| `reflect_max_tokens` | Max tokens for reflect | `max_tokens` |
|
||||||
|
| `reflect_response_schema` | JSON schema for reflect output | `None` |
|
||||||
|
| `reflect_tags` | Tags for reflect (fallback: `recall_tags`) | `None` |
|
||||||
|
| `reflect_tags_match` | Tag match for reflect | `recall_tags_match` |
|
||||||
|
|
||||||
|
## Advanced: Manual Registration
|
||||||
|
|
||||||
|
```python
|
||||||
|
from hindsight_ag2 import create_hindsight_tools
|
||||||
|
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="my-bank",
|
||||||
|
hindsight_api_url="http://localhost:8888",
|
||||||
|
)
|
||||||
|
for tool_fn in tools:
|
||||||
|
assistant.register_for_llm(description=tool_fn.__doc__)(tool_fn)
|
||||||
|
user_proxy.register_for_execution()(tool_fn)
|
||||||
|
```
|
||||||
|
|
||||||
|
## Advanced: GroupChat with Shared Memory
|
||||||
|
|
||||||
|
```python
|
||||||
|
from autogen import AssistantAgent, UserProxyAgent, GroupChat, GroupChatManager, LLMConfig
|
||||||
|
from hindsight_ag2 import register_hindsight_tools
|
||||||
|
|
||||||
|
llm_config = LLMConfig(api_type="openai", model="gpt-4o-mini")
|
||||||
|
|
||||||
|
with llm_config:
|
||||||
|
researcher = AssistantAgent(name="researcher", system_message="You research topics.")
|
||||||
|
writer = AssistantAgent(name="writer", system_message="You write content.")
|
||||||
|
executor = UserProxyAgent(name="executor", human_input_mode="NEVER")
|
||||||
|
|
||||||
|
# All agents share the same memory bank
|
||||||
|
for agent in [researcher, writer]:
|
||||||
|
register_hindsight_tools(agent, executor, bank_id="team-memory")
|
||||||
|
|
||||||
|
group_chat = GroupChat(agents=[researcher, writer, executor], messages=[])
|
||||||
|
manager = GroupChatManager(groupchat=group_chat)
|
||||||
|
```
|
||||||
|
|
||||||
|
## Requirements
|
||||||
|
|
||||||
|
- `ag2>=0.9.0`
|
||||||
|
- `hindsight-client>=0.4.0`
|
||||||
|
|
||||||
|
## Documentation
|
||||||
|
|
||||||
|
- [Hindsight Documentation](https://hindsight.docs.vectorize.io)
|
||||||
|
- [AG2 Documentation](https://docs.ag2.ai)
|
||||||
|
- [API Reference](https://hindsight.docs.vectorize.io/api)
|
||||||
48
hindsight-integrations/ag2/hindsight_ag2/__init__.py
Normal file
48
hindsight-integrations/ag2/hindsight_ag2/__init__.py
Normal file
|
|
@ -0,0 +1,48 @@
|
||||||
|
"""Hindsight-AG2: Persistent memory tools for AG2 agents.
|
||||||
|
|
||||||
|
Provides Hindsight-backed tool functions that give AG2 agents long-term
|
||||||
|
memory across conversations via retain/recall/reflect operations.
|
||||||
|
|
||||||
|
Basic usage::
|
||||||
|
|
||||||
|
from hindsight_ag2 import register_hindsight_tools
|
||||||
|
|
||||||
|
register_hindsight_tools(
|
||||||
|
assistant, user_proxy,
|
||||||
|
bank_id="user-123",
|
||||||
|
hindsight_api_url="http://localhost:8888",
|
||||||
|
)
|
||||||
|
|
||||||
|
Manual registration::
|
||||||
|
|
||||||
|
from hindsight_ag2 import create_hindsight_tools
|
||||||
|
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="user-123",
|
||||||
|
hindsight_api_url="http://localhost:8888",
|
||||||
|
)
|
||||||
|
for tool_fn in tools:
|
||||||
|
assistant.register_for_llm(description=tool_fn.__doc__)(tool_fn)
|
||||||
|
user_proxy.register_for_execution()(tool_fn)
|
||||||
|
"""
|
||||||
|
|
||||||
|
from .config import (
|
||||||
|
HindsightAG2Config,
|
||||||
|
configure,
|
||||||
|
get_config,
|
||||||
|
reset_config,
|
||||||
|
)
|
||||||
|
from .errors import HindsightError
|
||||||
|
from .tools import create_hindsight_tools, register_hindsight_tools
|
||||||
|
|
||||||
|
__version__ = "0.1.0"
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"configure",
|
||||||
|
"get_config",
|
||||||
|
"reset_config",
|
||||||
|
"HindsightAG2Config",
|
||||||
|
"HindsightError",
|
||||||
|
"create_hindsight_tools",
|
||||||
|
"register_hindsight_tools",
|
||||||
|
]
|
||||||
32
hindsight-integrations/ag2/hindsight_ag2/_client.py
Normal file
32
hindsight-integrations/ag2/hindsight_ag2/_client.py
Normal file
|
|
@ -0,0 +1,32 @@
|
||||||
|
"""Shared Hindsight client resolution logic."""
|
||||||
|
|
||||||
|
from typing import Any, Optional
|
||||||
|
|
||||||
|
from hindsight_client import Hindsight
|
||||||
|
|
||||||
|
from .config import get_config
|
||||||
|
from .errors import HindsightError
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_client(
|
||||||
|
client: Optional[Hindsight],
|
||||||
|
hindsight_api_url: Optional[str],
|
||||||
|
api_key: Optional[str],
|
||||||
|
) -> Hindsight:
|
||||||
|
"""Resolve a Hindsight client from explicit args or global config."""
|
||||||
|
if client is not None:
|
||||||
|
return client
|
||||||
|
|
||||||
|
config = get_config()
|
||||||
|
url = hindsight_api_url or (config.hindsight_api_url if config else None)
|
||||||
|
key = api_key or (config.api_key if config else None)
|
||||||
|
|
||||||
|
if url is None:
|
||||||
|
raise HindsightError(
|
||||||
|
"No Hindsight API URL configured. Pass client= or hindsight_api_url=, or call configure() first."
|
||||||
|
)
|
||||||
|
|
||||||
|
kwargs: dict[str, Any] = {"base_url": url, "timeout": 30.0}
|
||||||
|
if key:
|
||||||
|
kwargs["api_key"] = key
|
||||||
|
return Hindsight(**kwargs)
|
||||||
91
hindsight-integrations/ag2/hindsight_ag2/config.py
Normal file
91
hindsight-integrations/ag2/hindsight_ag2/config.py
Normal file
|
|
@ -0,0 +1,91 @@
|
||||||
|
"""Global configuration for Hindsight-AG2 integration."""
|
||||||
|
|
||||||
|
import os
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
DEFAULT_HINDSIGHT_API_URL = "https://api.hindsight.vectorize.io"
|
||||||
|
HINDSIGHT_API_KEY_ENV = "HINDSIGHT_API_KEY"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class HindsightAG2Config:
|
||||||
|
"""Connection and default settings for the AG2 integration.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
hindsight_api_url: URL of the Hindsight API server.
|
||||||
|
api_key: API key for Hindsight authentication.
|
||||||
|
budget: Default recall budget level (low/mid/high).
|
||||||
|
max_tokens: Default maximum tokens for recall results.
|
||||||
|
tags: Default tags applied when storing memories.
|
||||||
|
recall_tags: Default tags to filter when searching memories.
|
||||||
|
recall_tags_match: Tag matching mode (any/all/any_strict/all_strict).
|
||||||
|
verbose: Enable verbose logging.
|
||||||
|
"""
|
||||||
|
|
||||||
|
hindsight_api_url: str = DEFAULT_HINDSIGHT_API_URL
|
||||||
|
api_key: Optional[str] = None
|
||||||
|
budget: str = "mid"
|
||||||
|
max_tokens: int = 4096
|
||||||
|
tags: Optional[list[str]] = None
|
||||||
|
recall_tags: Optional[list[str]] = None
|
||||||
|
recall_tags_match: str = "any"
|
||||||
|
verbose: bool = False
|
||||||
|
|
||||||
|
|
||||||
|
_global_config: Optional[HindsightAG2Config] = None
|
||||||
|
|
||||||
|
|
||||||
|
def configure(
|
||||||
|
hindsight_api_url: Optional[str] = None,
|
||||||
|
api_key: Optional[str] = None,
|
||||||
|
budget: str = "mid",
|
||||||
|
max_tokens: int = 4096,
|
||||||
|
tags: Optional[list[str]] = None,
|
||||||
|
recall_tags: Optional[list[str]] = None,
|
||||||
|
recall_tags_match: str = "any",
|
||||||
|
verbose: bool = False,
|
||||||
|
) -> HindsightAG2Config:
|
||||||
|
"""Configure Hindsight connection and default settings.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
hindsight_api_url: Hindsight API URL (default: production).
|
||||||
|
api_key: API key. Falls back to HINDSIGHT_API_KEY env var.
|
||||||
|
budget: Default recall budget (low/mid/high).
|
||||||
|
max_tokens: Default max tokens for recall.
|
||||||
|
tags: Default tags for retain operations.
|
||||||
|
recall_tags: Default tags to filter recall/search.
|
||||||
|
recall_tags_match: Tag matching mode.
|
||||||
|
verbose: Enable verbose logging.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
The configured HindsightAG2Config.
|
||||||
|
"""
|
||||||
|
global _global_config
|
||||||
|
|
||||||
|
resolved_url = hindsight_api_url or DEFAULT_HINDSIGHT_API_URL
|
||||||
|
resolved_key = api_key or os.environ.get(HINDSIGHT_API_KEY_ENV)
|
||||||
|
|
||||||
|
_global_config = HindsightAG2Config(
|
||||||
|
hindsight_api_url=resolved_url,
|
||||||
|
api_key=resolved_key,
|
||||||
|
budget=budget,
|
||||||
|
max_tokens=max_tokens,
|
||||||
|
tags=tags,
|
||||||
|
recall_tags=recall_tags,
|
||||||
|
recall_tags_match=recall_tags_match,
|
||||||
|
verbose=verbose,
|
||||||
|
)
|
||||||
|
|
||||||
|
return _global_config
|
||||||
|
|
||||||
|
|
||||||
|
def get_config() -> Optional[HindsightAG2Config]:
|
||||||
|
"""Get the current global configuration."""
|
||||||
|
return _global_config
|
||||||
|
|
||||||
|
|
||||||
|
def reset_config() -> None:
|
||||||
|
"""Reset global configuration to None."""
|
||||||
|
global _global_config
|
||||||
|
_global_config = None
|
||||||
7
hindsight-integrations/ag2/hindsight_ag2/errors.py
Normal file
7
hindsight-integrations/ag2/hindsight_ag2/errors.py
Normal file
|
|
@ -0,0 +1,7 @@
|
||||||
|
"""Hindsight-AG2 error types."""
|
||||||
|
|
||||||
|
|
||||||
|
class HindsightError(Exception):
|
||||||
|
"""Exception raised when a Hindsight memory operation fails."""
|
||||||
|
|
||||||
|
pass
|
||||||
265
hindsight-integrations/ag2/hindsight_ag2/tools.py
Normal file
265
hindsight-integrations/ag2/hindsight_ag2/tools.py
Normal file
|
|
@ -0,0 +1,265 @@
|
||||||
|
"""AG2 tool definitions for Hindsight memory operations.
|
||||||
|
|
||||||
|
Provides factory functions that create AG2-compatible tool functions
|
||||||
|
backed by Hindsight's retain/recall/reflect APIs. Tools are plain Python
|
||||||
|
functions with ``Annotated`` type hints, compatible with AG2's
|
||||||
|
``@register_for_llm`` / ``@register_for_execution`` pattern.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from typing import Annotated, Any, Optional
|
||||||
|
|
||||||
|
from hindsight_client import Hindsight
|
||||||
|
|
||||||
|
from ._client import resolve_client
|
||||||
|
from .config import get_config
|
||||||
|
from .errors import HindsightError
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def create_hindsight_tools(
|
||||||
|
*,
|
||||||
|
bank_id: str,
|
||||||
|
client: Optional[Hindsight] = None,
|
||||||
|
hindsight_api_url: Optional[str] = None,
|
||||||
|
api_key: Optional[str] = None,
|
||||||
|
budget: Optional[str] = None,
|
||||||
|
max_tokens: Optional[int] = None,
|
||||||
|
tags: Optional[list[str]] = None,
|
||||||
|
recall_tags: Optional[list[str]] = None,
|
||||||
|
recall_tags_match: Optional[str] = None,
|
||||||
|
# Retain options
|
||||||
|
retain_metadata: Optional[dict[str, str]] = None,
|
||||||
|
retain_document_id: Optional[str] = None,
|
||||||
|
# Recall options
|
||||||
|
recall_types: Optional[list[str]] = None,
|
||||||
|
recall_include_entities: bool = False,
|
||||||
|
# Reflect options
|
||||||
|
reflect_context: Optional[str] = None,
|
||||||
|
reflect_max_tokens: Optional[int] = None,
|
||||||
|
reflect_response_schema: Optional[dict[str, Any]] = None,
|
||||||
|
reflect_tags: Optional[list[str]] = None,
|
||||||
|
reflect_tags_match: Optional[str] = None,
|
||||||
|
include_retain: bool = True,
|
||||||
|
include_recall: bool = True,
|
||||||
|
include_reflect: bool = True,
|
||||||
|
) -> list:
|
||||||
|
"""Create Hindsight memory tools for AG2 agents.
|
||||||
|
|
||||||
|
Returns a list of plain Python functions compatible with AG2's
|
||||||
|
``@register_for_llm`` / ``@register_for_execution`` pattern.
|
||||||
|
Each function uses ``Annotated`` type hints for parameter descriptions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
bank_id: The Hindsight memory bank to operate on.
|
||||||
|
client: Pre-configured Hindsight client (preferred).
|
||||||
|
hindsight_api_url: API URL (used if no client provided).
|
||||||
|
api_key: API key (used if no client provided).
|
||||||
|
budget: Recall/reflect budget level (low/mid/high).
|
||||||
|
max_tokens: Maximum tokens for recall results.
|
||||||
|
tags: Tags applied when storing memories via retain.
|
||||||
|
recall_tags: Tags to filter when searching memories.
|
||||||
|
recall_tags_match: Tag matching mode (any/all/any_strict/all_strict).
|
||||||
|
retain_metadata: Default metadata dict for retain operations.
|
||||||
|
retain_document_id: Default document_id for retain (groups/upserts memories).
|
||||||
|
recall_types: Fact types to filter (world, experience, opinion, observation).
|
||||||
|
recall_include_entities: Include entity information in recall results.
|
||||||
|
reflect_context: Additional context for reflect operations.
|
||||||
|
reflect_max_tokens: Max tokens for reflect results (defaults to max_tokens).
|
||||||
|
reflect_response_schema: JSON schema to constrain reflect output format.
|
||||||
|
reflect_tags: Tags to filter memories used in reflect (defaults to recall_tags).
|
||||||
|
reflect_tags_match: Tag matching for reflect (defaults to recall_tags_match).
|
||||||
|
include_retain: Include the retain (store) tool.
|
||||||
|
include_recall: Include the recall (search) tool.
|
||||||
|
include_reflect: Include the reflect (synthesize) tool.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of callable tool functions.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
HindsightError: If no client or API URL can be resolved.
|
||||||
|
|
||||||
|
Usage::
|
||||||
|
|
||||||
|
tools = create_hindsight_tools(bank_id="my-bank", client=client)
|
||||||
|
for tool_fn in tools:
|
||||||
|
agent.register_for_llm(description=tool_fn.__doc__)(tool_fn)
|
||||||
|
executor.register_for_execution()(tool_fn)
|
||||||
|
"""
|
||||||
|
resolved_client = resolve_client(client, hindsight_api_url, api_key)
|
||||||
|
|
||||||
|
config = get_config()
|
||||||
|
effective_tags = tags if tags is not None else (config.tags if config else None)
|
||||||
|
effective_recall_tags = (
|
||||||
|
recall_tags
|
||||||
|
if recall_tags is not None
|
||||||
|
else (config.recall_tags if config else None)
|
||||||
|
)
|
||||||
|
effective_recall_tags_match = (
|
||||||
|
recall_tags_match
|
||||||
|
if recall_tags_match is not None
|
||||||
|
else (config.recall_tags_match if config else "any")
|
||||||
|
)
|
||||||
|
effective_budget = (
|
||||||
|
budget if budget is not None else (config.budget if config else "mid")
|
||||||
|
)
|
||||||
|
effective_max_tokens = (
|
||||||
|
max_tokens
|
||||||
|
if max_tokens is not None
|
||||||
|
else (config.max_tokens if config else 4096)
|
||||||
|
)
|
||||||
|
|
||||||
|
tools: list = []
|
||||||
|
|
||||||
|
if include_retain:
|
||||||
|
|
||||||
|
def hindsight_retain(
|
||||||
|
content: Annotated[
|
||||||
|
str,
|
||||||
|
"The information to store in long-term memory. Include important facts, "
|
||||||
|
"user preferences, decisions, or anything that should be remembered across conversations.",
|
||||||
|
],
|
||||||
|
) -> str:
|
||||||
|
"""Store information to long-term memory for later retrieval.
|
||||||
|
|
||||||
|
Use this to save important facts, user preferences, decisions,
|
||||||
|
or any information that should be remembered across conversations.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
retain_kwargs: dict[str, Any] = {"bank_id": bank_id, "content": content}
|
||||||
|
if effective_tags:
|
||||||
|
retain_kwargs["tags"] = effective_tags
|
||||||
|
if retain_metadata:
|
||||||
|
retain_kwargs["metadata"] = retain_metadata
|
||||||
|
if retain_document_id:
|
||||||
|
retain_kwargs["document_id"] = retain_document_id
|
||||||
|
resolved_client.retain(**retain_kwargs)
|
||||||
|
return "Memory stored successfully."
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Retain failed: {e}")
|
||||||
|
raise HindsightError(f"Retain failed: {e}") from e
|
||||||
|
|
||||||
|
tools.append(hindsight_retain)
|
||||||
|
|
||||||
|
if include_recall:
|
||||||
|
|
||||||
|
def hindsight_recall(
|
||||||
|
query: Annotated[
|
||||||
|
str,
|
||||||
|
"The search query to find relevant memories. Be specific about what information you're looking for.",
|
||||||
|
],
|
||||||
|
) -> str:
|
||||||
|
"""Search long-term memory for relevant information.
|
||||||
|
|
||||||
|
Use this to find previously stored facts, preferences, or context.
|
||||||
|
Returns a numbered list of matching memories.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
recall_kwargs: dict[str, Any] = {
|
||||||
|
"bank_id": bank_id,
|
||||||
|
"query": query,
|
||||||
|
"budget": effective_budget,
|
||||||
|
"max_tokens": effective_max_tokens,
|
||||||
|
}
|
||||||
|
if effective_recall_tags:
|
||||||
|
recall_kwargs["tags"] = effective_recall_tags
|
||||||
|
recall_kwargs["tags_match"] = effective_recall_tags_match
|
||||||
|
if recall_types:
|
||||||
|
recall_kwargs["types"] = recall_types
|
||||||
|
if recall_include_entities:
|
||||||
|
recall_kwargs["include_entities"] = True
|
||||||
|
response = resolved_client.recall(**recall_kwargs)
|
||||||
|
if not response.results:
|
||||||
|
return "No relevant memories found."
|
||||||
|
lines = []
|
||||||
|
for i, result in enumerate(response.results, 1):
|
||||||
|
lines.append(f"{i}. {result.text}")
|
||||||
|
return "\n".join(lines)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Recall failed: {e}")
|
||||||
|
raise HindsightError(f"Recall failed: {e}") from e
|
||||||
|
|
||||||
|
tools.append(hindsight_recall)
|
||||||
|
|
||||||
|
if include_reflect:
|
||||||
|
|
||||||
|
def hindsight_reflect(
|
||||||
|
query: Annotated[
|
||||||
|
str,
|
||||||
|
"The question or topic to synthesize a thoughtful answer about from long-term memories.",
|
||||||
|
],
|
||||||
|
) -> str:
|
||||||
|
"""Synthesize a thoughtful answer from long-term memories.
|
||||||
|
|
||||||
|
Use this when you need a coherent summary or reasoned response
|
||||||
|
about what you know, rather than raw memory facts.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
reflect_kwargs: dict[str, Any] = {
|
||||||
|
"bank_id": bank_id,
|
||||||
|
"query": query,
|
||||||
|
"budget": effective_budget,
|
||||||
|
}
|
||||||
|
if reflect_context:
|
||||||
|
reflect_kwargs["context"] = reflect_context
|
||||||
|
effective_reflect_max = reflect_max_tokens or effective_max_tokens
|
||||||
|
if effective_reflect_max:
|
||||||
|
reflect_kwargs["max_tokens"] = effective_reflect_max
|
||||||
|
if reflect_response_schema:
|
||||||
|
reflect_kwargs["response_schema"] = reflect_response_schema
|
||||||
|
# Reflect tags: use reflect-specific or fall back to recall tags
|
||||||
|
effective_reflect_tags = (
|
||||||
|
reflect_tags if reflect_tags is not None else effective_recall_tags
|
||||||
|
)
|
||||||
|
effective_reflect_tags_match = (
|
||||||
|
reflect_tags_match or effective_recall_tags_match
|
||||||
|
)
|
||||||
|
if effective_reflect_tags:
|
||||||
|
reflect_kwargs["tags"] = effective_reflect_tags
|
||||||
|
reflect_kwargs["tags_match"] = effective_reflect_tags_match
|
||||||
|
response = resolved_client.reflect(**reflect_kwargs)
|
||||||
|
return response.text or "No relevant memories found."
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Reflect failed: {e}")
|
||||||
|
raise HindsightError(f"Reflect failed: {e}") from e
|
||||||
|
|
||||||
|
tools.append(hindsight_reflect)
|
||||||
|
|
||||||
|
return tools
|
||||||
|
|
||||||
|
|
||||||
|
def register_hindsight_tools(
|
||||||
|
agent,
|
||||||
|
executor,
|
||||||
|
*,
|
||||||
|
bank_id: str,
|
||||||
|
**kwargs,
|
||||||
|
) -> list:
|
||||||
|
"""Convenience: create tools AND register them on AG2 agents.
|
||||||
|
|
||||||
|
Creates Hindsight memory tools and registers them on the given AG2
|
||||||
|
agents using ``register_for_llm`` and ``register_for_execution``.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
agent: AG2 agent to register tools for LLM calling.
|
||||||
|
executor: AG2 agent to register tools for execution.
|
||||||
|
bank_id: Hindsight memory bank ID.
|
||||||
|
**kwargs: All other args passed to ``create_hindsight_tools()``.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of registered tool functions.
|
||||||
|
|
||||||
|
Usage::
|
||||||
|
|
||||||
|
tools = register_hindsight_tools(
|
||||||
|
assistant, user_proxy,
|
||||||
|
bank_id="my-bank",
|
||||||
|
hindsight_api_url="http://localhost:8888",
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
tools = create_hindsight_tools(bank_id=bank_id, **kwargs)
|
||||||
|
for tool_fn in tools:
|
||||||
|
agent.register_for_llm(description=tool_fn.__doc__)(tool_fn)
|
||||||
|
executor.register_for_execution()(tool_fn)
|
||||||
|
return tools
|
||||||
54
hindsight-integrations/ag2/pyproject.toml
Normal file
54
hindsight-integrations/ag2/pyproject.toml
Normal file
|
|
@ -0,0 +1,54 @@
|
||||||
|
[project]
|
||||||
|
name = "hindsight-ag2"
|
||||||
|
version = "0.1.0"
|
||||||
|
description = "AG2 integration for Hindsight - persistent memory tools for AI agents"
|
||||||
|
readme = "README.md"
|
||||||
|
requires-python = ">=3.10"
|
||||||
|
license = { text = "MIT" }
|
||||||
|
authors = [
|
||||||
|
{ name = "Vectorize", email = "support@vectorize.io" }
|
||||||
|
]
|
||||||
|
keywords = [
|
||||||
|
"ai",
|
||||||
|
"memory",
|
||||||
|
"ag2",
|
||||||
|
"autogen",
|
||||||
|
"agents",
|
||||||
|
"hindsight",
|
||||||
|
]
|
||||||
|
classifiers = [
|
||||||
|
"Development Status :: 4 - Beta",
|
||||||
|
"Intended Audience :: Developers",
|
||||||
|
"License :: OSI Approved :: MIT License",
|
||||||
|
"Programming Language :: Python :: 3",
|
||||||
|
"Programming Language :: Python :: 3.10",
|
||||||
|
"Programming Language :: Python :: 3.11",
|
||||||
|
"Programming Language :: Python :: 3.12",
|
||||||
|
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||||
|
]
|
||||||
|
|
||||||
|
dependencies = [
|
||||||
|
"ag2>=0.9.0",
|
||||||
|
"hindsight-client>=0.4.0",
|
||||||
|
]
|
||||||
|
|
||||||
|
[project.urls]
|
||||||
|
Homepage = "https://github.com/vectorize-io/hindsight"
|
||||||
|
Documentation = "https://github.com/vectorize-io/hindsight/tree/main/hindsight-integrations/ag2"
|
||||||
|
Repository = "https://github.com/vectorize-io/hindsight"
|
||||||
|
|
||||||
|
[build-system]
|
||||||
|
requires = ["hatchling"]
|
||||||
|
build-backend = "hatchling.build"
|
||||||
|
|
||||||
|
[tool.hatch.build.targets.wheel]
|
||||||
|
packages = ["hindsight_ag2"]
|
||||||
|
|
||||||
|
[tool.pytest.ini_options]
|
||||||
|
testpaths = ["tests"]
|
||||||
|
|
||||||
|
[dependency-groups]
|
||||||
|
dev = [
|
||||||
|
"pytest>=9.0.2",
|
||||||
|
"pytest-mock>=3.10.0",
|
||||||
|
]
|
||||||
0
hindsight-integrations/ag2/tests/__init__.py
Normal file
0
hindsight-integrations/ag2/tests/__init__.py
Normal file
580
hindsight-integrations/ag2/tests/test_tools.py
Normal file
580
hindsight-integrations/ag2/tests/test_tools.py
Normal file
|
|
@ -0,0 +1,580 @@
|
||||||
|
"""Unit tests for Hindsight AG2 tools."""
|
||||||
|
|
||||||
|
import inspect
|
||||||
|
from typing import Annotated, get_type_hints
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from hindsight_ag2 import (
|
||||||
|
configure,
|
||||||
|
create_hindsight_tools,
|
||||||
|
register_hindsight_tools,
|
||||||
|
reset_config,
|
||||||
|
)
|
||||||
|
from hindsight_ag2.errors import HindsightError
|
||||||
|
|
||||||
|
|
||||||
|
def _mock_client():
|
||||||
|
"""Create a mock Hindsight client with sync methods."""
|
||||||
|
client = MagicMock()
|
||||||
|
client.retain = MagicMock(return_value=None)
|
||||||
|
client.recall = MagicMock()
|
||||||
|
client.reflect = MagicMock()
|
||||||
|
return client
|
||||||
|
|
||||||
|
|
||||||
|
def _mock_recall_response(texts: list[str]):
|
||||||
|
response = MagicMock()
|
||||||
|
results = []
|
||||||
|
for t in texts:
|
||||||
|
r = MagicMock()
|
||||||
|
r.text = t
|
||||||
|
results.append(r)
|
||||||
|
response.results = results
|
||||||
|
return response
|
||||||
|
|
||||||
|
|
||||||
|
def _mock_reflect_response(text: str):
|
||||||
|
response = MagicMock()
|
||||||
|
response.text = text
|
||||||
|
return response
|
||||||
|
|
||||||
|
|
||||||
|
class TestImports:
|
||||||
|
def test_imports(self):
|
||||||
|
from hindsight_ag2 import ( # noqa: F401
|
||||||
|
HindsightAG2Config,
|
||||||
|
HindsightError,
|
||||||
|
configure,
|
||||||
|
create_hindsight_tools,
|
||||||
|
get_config,
|
||||||
|
register_hindsight_tools,
|
||||||
|
reset_config,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestCreateHindsightTools:
|
||||||
|
def setup_method(self):
|
||||||
|
reset_config()
|
||||||
|
|
||||||
|
def teardown_method(self):
|
||||||
|
reset_config()
|
||||||
|
|
||||||
|
def test_returns_three_tools_by_default(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(bank_id="test", client=client)
|
||||||
|
assert len(tools) == 3
|
||||||
|
|
||||||
|
def test_include_retain_only(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=True,
|
||||||
|
include_recall=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
assert len(tools) == 1
|
||||||
|
assert tools[0].__name__ == "hindsight_retain"
|
||||||
|
|
||||||
|
def test_include_recall_only(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_recall=True,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
assert len(tools) == 1
|
||||||
|
assert tools[0].__name__ == "hindsight_recall"
|
||||||
|
|
||||||
|
def test_include_reflect_only(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_recall=False,
|
||||||
|
include_reflect=True,
|
||||||
|
)
|
||||||
|
assert len(tools) == 1
|
||||||
|
assert tools[0].__name__ == "hindsight_reflect"
|
||||||
|
|
||||||
|
def test_no_tools_when_all_excluded(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_recall=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
assert len(tools) == 0
|
||||||
|
|
||||||
|
def test_tool_names(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(bank_id="test", client=client)
|
||||||
|
names = [fn.__name__ for fn in tools]
|
||||||
|
assert names == ["hindsight_retain", "hindsight_recall", "hindsight_reflect"]
|
||||||
|
|
||||||
|
def test_tool_docstrings(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(bank_id="test", client=client)
|
||||||
|
for fn in tools:
|
||||||
|
assert fn.__doc__ is not None
|
||||||
|
assert len(fn.__doc__) > 0
|
||||||
|
|
||||||
|
def test_raises_without_client_or_config(self):
|
||||||
|
with pytest.raises(HindsightError, match="No Hindsight API URL"):
|
||||||
|
create_hindsight_tools(bank_id="test")
|
||||||
|
|
||||||
|
def test_falls_back_to_global_config(self):
|
||||||
|
configure(hindsight_api_url="http://localhost:8888")
|
||||||
|
with patch("hindsight_ag2._client.Hindsight") as mock_cls:
|
||||||
|
mock_cls.return_value = _mock_client()
|
||||||
|
tools = create_hindsight_tools(bank_id="test")
|
||||||
|
assert len(tools) == 3
|
||||||
|
mock_cls.assert_called_once_with(
|
||||||
|
base_url="http://localhost:8888", timeout=30.0
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_explicit_url_overrides_config(self):
|
||||||
|
configure(hindsight_api_url="http://config:8888")
|
||||||
|
with patch("hindsight_ag2._client.Hindsight") as mock_cls:
|
||||||
|
mock_cls.return_value = _mock_client()
|
||||||
|
create_hindsight_tools(
|
||||||
|
bank_id="test", hindsight_api_url="http://explicit:9999"
|
||||||
|
)
|
||||||
|
mock_cls.assert_called_once_with(
|
||||||
|
base_url="http://explicit:9999", timeout=30.0
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestRetainTool:
|
||||||
|
def test_retain_stores_memory(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test-bank",
|
||||||
|
client=client,
|
||||||
|
include_recall=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
result = tools[0]("The user likes Python")
|
||||||
|
assert result == "Memory stored successfully."
|
||||||
|
client.retain.assert_called_once_with(
|
||||||
|
bank_id="test-bank", content="The user likes Python"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_retain_passes_tags(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test-bank",
|
||||||
|
client=client,
|
||||||
|
tags=["source:chat"],
|
||||||
|
include_recall=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
tools[0]("some content")
|
||||||
|
call_kwargs = client.retain.call_args[1]
|
||||||
|
assert call_kwargs["tags"] == ["source:chat"]
|
||||||
|
|
||||||
|
def test_retain_passes_metadata(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
retain_metadata={"source": "chat", "session": "abc"},
|
||||||
|
include_recall=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
tools[0]("content")
|
||||||
|
call_kwargs = client.retain.call_args[1]
|
||||||
|
assert call_kwargs["metadata"] == {"source": "chat", "session": "abc"}
|
||||||
|
|
||||||
|
def test_retain_passes_document_id(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
retain_document_id="session-123",
|
||||||
|
include_recall=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
tools[0]("content")
|
||||||
|
call_kwargs = client.retain.call_args[1]
|
||||||
|
assert call_kwargs["document_id"] == "session-123"
|
||||||
|
|
||||||
|
def test_retain_raises_hindsight_error(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.retain.side_effect = RuntimeError("connection refused")
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_recall=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
with pytest.raises(HindsightError, match="Retain failed"):
|
||||||
|
tools[0]("content")
|
||||||
|
|
||||||
|
|
||||||
|
class TestRecallTool:
|
||||||
|
def test_recall_returns_numbered_results(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.recall.return_value = _mock_recall_response(
|
||||||
|
["User likes Python", "User is in NYC"]
|
||||||
|
)
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test-bank",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
result = tools[0]("user preferences")
|
||||||
|
assert "1. User likes Python" in result
|
||||||
|
assert "2. User is in NYC" in result
|
||||||
|
|
||||||
|
def test_recall_empty_results(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.recall.return_value = _mock_recall_response([])
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
result = tools[0]("anything")
|
||||||
|
assert result == "No relevant memories found."
|
||||||
|
|
||||||
|
def test_recall_passes_budget_and_max_tokens(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.recall.return_value = _mock_recall_response(["fact"])
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
budget="high",
|
||||||
|
max_tokens=2048,
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = client.recall.call_args[1]
|
||||||
|
assert call_kwargs["budget"] == "high"
|
||||||
|
assert call_kwargs["max_tokens"] == 2048
|
||||||
|
|
||||||
|
def test_recall_passes_tags(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.recall.return_value = _mock_recall_response(["fact"])
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
recall_tags=["scope:user"],
|
||||||
|
recall_tags_match="all",
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = client.recall.call_args[1]
|
||||||
|
assert call_kwargs["tags"] == ["scope:user"]
|
||||||
|
assert call_kwargs["tags_match"] == "all"
|
||||||
|
|
||||||
|
def test_recall_passes_types(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.recall.return_value = _mock_recall_response(["fact"])
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
recall_types=["world", "experience"],
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = client.recall.call_args[1]
|
||||||
|
assert call_kwargs["types"] == ["world", "experience"]
|
||||||
|
|
||||||
|
def test_recall_passes_include_entities(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.recall.return_value = _mock_recall_response(["fact"])
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
recall_include_entities=True,
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = client.recall.call_args[1]
|
||||||
|
assert call_kwargs["include_entities"] is True
|
||||||
|
|
||||||
|
def test_recall_raises_hindsight_error(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.recall.side_effect = RuntimeError("connection refused")
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
with pytest.raises(HindsightError, match="Recall failed"):
|
||||||
|
tools[0]("query")
|
||||||
|
|
||||||
|
|
||||||
|
class TestReflectTool:
|
||||||
|
def test_reflect_returns_text(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.reflect.return_value = _mock_reflect_response(
|
||||||
|
"The user is a Python developer who prefers functional patterns."
|
||||||
|
)
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test-bank",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_recall=False,
|
||||||
|
)
|
||||||
|
result = tools[0]("What do you know about the user?")
|
||||||
|
assert result == "The user is a Python developer who prefers functional patterns."
|
||||||
|
|
||||||
|
def test_reflect_empty_returns_fallback(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.reflect.return_value = _mock_reflect_response("")
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_recall=False,
|
||||||
|
)
|
||||||
|
result = tools[0]("anything")
|
||||||
|
assert result == "No relevant memories found."
|
||||||
|
|
||||||
|
def test_reflect_passes_budget(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.reflect.return_value = _mock_reflect_response("answer")
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
budget="high",
|
||||||
|
include_retain=False,
|
||||||
|
include_recall=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = client.reflect.call_args[1]
|
||||||
|
assert call_kwargs["budget"] == "high"
|
||||||
|
|
||||||
|
def test_reflect_passes_context(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.reflect.return_value = _mock_reflect_response("answer")
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
reflect_context="The user is asking about project setup",
|
||||||
|
include_retain=False,
|
||||||
|
include_recall=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = client.reflect.call_args[1]
|
||||||
|
assert call_kwargs["context"] == "The user is asking about project setup"
|
||||||
|
|
||||||
|
def test_reflect_passes_max_tokens_and_response_schema(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.reflect.return_value = _mock_reflect_response("answer")
|
||||||
|
schema = {"type": "object", "properties": {"summary": {"type": "string"}}}
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
reflect_max_tokens=2048,
|
||||||
|
reflect_response_schema=schema,
|
||||||
|
include_retain=False,
|
||||||
|
include_recall=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = client.reflect.call_args[1]
|
||||||
|
assert call_kwargs["max_tokens"] == 2048
|
||||||
|
assert call_kwargs["response_schema"] == schema
|
||||||
|
|
||||||
|
def test_reflect_passes_tags(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.reflect.return_value = _mock_reflect_response("answer")
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
reflect_tags=["scope:global"],
|
||||||
|
reflect_tags_match="all",
|
||||||
|
include_retain=False,
|
||||||
|
include_recall=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = client.reflect.call_args[1]
|
||||||
|
assert call_kwargs["tags"] == ["scope:global"]
|
||||||
|
assert call_kwargs["tags_match"] == "all"
|
||||||
|
|
||||||
|
def test_reflect_raises_hindsight_error(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.reflect.side_effect = RuntimeError("connection refused")
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_recall=False,
|
||||||
|
)
|
||||||
|
with pytest.raises(HindsightError, match="Reflect failed"):
|
||||||
|
tools[0]("query")
|
||||||
|
|
||||||
|
|
||||||
|
class TestAnnotatedTypes:
|
||||||
|
def test_retain_has_annotated_parameter(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_recall=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
hints = get_type_hints(tools[0], include_extras=True)
|
||||||
|
assert "content" in hints
|
||||||
|
# Check it's Annotated
|
||||||
|
assert hasattr(hints["content"], "__metadata__")
|
||||||
|
|
||||||
|
def test_recall_has_annotated_parameter(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
hints = get_type_hints(tools[0], include_extras=True)
|
||||||
|
assert "query" in hints
|
||||||
|
assert hasattr(hints["query"], "__metadata__")
|
||||||
|
|
||||||
|
def test_reflect_has_annotated_parameter(self):
|
||||||
|
client = _mock_client()
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_recall=False,
|
||||||
|
)
|
||||||
|
hints = get_type_hints(tools[0], include_extras=True)
|
||||||
|
assert "query" in hints
|
||||||
|
assert hasattr(hints["query"], "__metadata__")
|
||||||
|
|
||||||
|
|
||||||
|
class TestRegisterHindsightTools:
|
||||||
|
def test_registers_all_tools(self):
|
||||||
|
client = _mock_client()
|
||||||
|
agent = MagicMock()
|
||||||
|
executor = MagicMock()
|
||||||
|
|
||||||
|
# Make register_for_llm return a decorator that returns the function
|
||||||
|
agent.register_for_llm.return_value = lambda fn: fn
|
||||||
|
executor.register_for_execution.return_value = lambda fn: fn
|
||||||
|
|
||||||
|
tools = register_hindsight_tools(
|
||||||
|
agent, executor, bank_id="test", client=client
|
||||||
|
)
|
||||||
|
|
||||||
|
assert len(tools) == 3
|
||||||
|
assert agent.register_for_llm.call_count == 3
|
||||||
|
assert executor.register_for_execution.call_count == 3
|
||||||
|
|
||||||
|
def test_registers_with_docstring_descriptions(self):
|
||||||
|
client = _mock_client()
|
||||||
|
agent = MagicMock()
|
||||||
|
executor = MagicMock()
|
||||||
|
|
||||||
|
agent.register_for_llm.return_value = lambda fn: fn
|
||||||
|
executor.register_for_execution.return_value = lambda fn: fn
|
||||||
|
|
||||||
|
register_hindsight_tools(agent, executor, bank_id="test", client=client)
|
||||||
|
|
||||||
|
# Each register_for_llm call should have a description kwarg
|
||||||
|
for call in agent.register_for_llm.call_args_list:
|
||||||
|
assert "description" in call.kwargs
|
||||||
|
assert call.kwargs["description"] is not None
|
||||||
|
assert len(call.kwargs["description"]) > 0
|
||||||
|
|
||||||
|
def test_passes_kwargs_to_create_tools(self):
|
||||||
|
client = _mock_client()
|
||||||
|
agent = MagicMock()
|
||||||
|
executor = MagicMock()
|
||||||
|
|
||||||
|
agent.register_for_llm.return_value = lambda fn: fn
|
||||||
|
executor.register_for_execution.return_value = lambda fn: fn
|
||||||
|
|
||||||
|
tools = register_hindsight_tools(
|
||||||
|
agent,
|
||||||
|
executor,
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=True,
|
||||||
|
include_recall=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert len(tools) == 1
|
||||||
|
assert tools[0].__name__ == "hindsight_retain"
|
||||||
|
|
||||||
|
|
||||||
|
class TestConfigDefaults:
|
||||||
|
def setup_method(self):
|
||||||
|
reset_config()
|
||||||
|
|
||||||
|
def teardown_method(self):
|
||||||
|
reset_config()
|
||||||
|
|
||||||
|
def test_default_budget(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.recall.return_value = _mock_recall_response(["fact"])
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = client.recall.call_args[1]
|
||||||
|
assert call_kwargs["budget"] == "mid"
|
||||||
|
|
||||||
|
def test_default_max_tokens(self):
|
||||||
|
client = _mock_client()
|
||||||
|
client.recall.return_value = _mock_recall_response(["fact"])
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
client=client,
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = client.recall.call_args[1]
|
||||||
|
assert call_kwargs["max_tokens"] == 4096
|
||||||
|
|
||||||
|
def test_config_budget_used_when_no_explicit(self):
|
||||||
|
configure(hindsight_api_url="http://localhost:8888", budget="low")
|
||||||
|
with patch("hindsight_ag2._client.Hindsight") as mock_cls:
|
||||||
|
mock_client = _mock_client()
|
||||||
|
mock_client.recall.return_value = _mock_recall_response(["fact"])
|
||||||
|
mock_cls.return_value = mock_client
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = mock_client.recall.call_args[1]
|
||||||
|
assert call_kwargs["budget"] == "low"
|
||||||
|
|
||||||
|
def test_explicit_budget_overrides_config(self):
|
||||||
|
configure(hindsight_api_url="http://localhost:8888", budget="low")
|
||||||
|
with patch("hindsight_ag2._client.Hindsight") as mock_cls:
|
||||||
|
mock_client = _mock_client()
|
||||||
|
mock_client.recall.return_value = _mock_recall_response(["fact"])
|
||||||
|
mock_cls.return_value = mock_client
|
||||||
|
tools = create_hindsight_tools(
|
||||||
|
bank_id="test",
|
||||||
|
budget="high",
|
||||||
|
include_retain=False,
|
||||||
|
include_reflect=False,
|
||||||
|
)
|
||||||
|
tools[0]("query")
|
||||||
|
call_kwargs = mock_client.recall.call_args[1]
|
||||||
|
assert call_kwargs["budget"] == "high"
|
||||||
1539
hindsight-integrations/ag2/uv.lock
Normal file
1539
hindsight-integrations/ag2/uv.lock
Normal file
File diff suppressed because it is too large
Load diff
Loading…
Reference in a new issue