holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

gemini support

Browse source
This commit is contained in:
Nicolò Boschi 2025-12-05 01:00:25 +01:00
parent ebe468e54f
commit 5ad2dfe03e
6 changed files with 293 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

145
hindsight-api/hindsight_api/engine/llm_wrapper.py
View file

@ -6,6 +6,9 @@ import time
import asyncio
from typing import Optional, Any, Dict, List
from openai import AsyncOpenAI, RateLimitError, APIError, APIStatusError, LengthFinishReasonError
from google import genai
from google.genai import types as genai_types
from google.genai import errors as genai_errors
import logging
logger = logging.getLogger(__name__)
@ -53,9 +56,9 @@ class LLMConfig:
self.model = model
# Validate provider
if self.provider not in ["openai", "groq", "ollama"]:
if self.provider not in ["openai", "groq", "ollama", "gemini"]:
raise ValueError(
f"Invalid LLM provider: {self.provider}. Must be 'openai', 'groq', or 'ollama'."
f"Invalid LLM provider: {self.provider}. Must be 'openai', 'groq', 'ollama', or 'gemini'."
)
# Set default base URLs
@ -66,19 +69,25 @@ class LLMConfig:
self.base_url = "http://localhost:11434/v1"
# Validate API key (not needed for ollama)
if self.provider != "ollama" and not self.api_key:
if self.provider not in ["ollama"] and not self.api_key:
raise ValueError(
f"API key not found for {self.provider}"
)
# Create client (private - use .call() method instead)
# Disable automatic retries - we handle retries in the call() method
if self.provider == "ollama":
if self.provider == "gemini":
self._gemini_client = genai.Client(api_key=self.api_key)
self._client = None # Not used for Gemini
elif self.provider == "ollama":
self._client = AsyncOpenAI(api_key="ollama", base_url=self.base_url, max_retries=0)
self._gemini_client = None
elif self.base_url:
self._client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url, max_retries=0)
self._gemini_client = None
else:
self._client = AsyncOpenAI(api_key=self.api_key, max_retries=0)
self._gemini_client = None
logger.info(
f"Initialized LLM: provider={self.provider}, model={self.model}, base_url={self.base_url}"
@ -116,6 +125,11 @@ class LLMConfig:
# Use global semaphore to limit concurrent requests
async with _global_llm_semaphore:
start_time = time.time()
import json
# Handle Gemini provider separately
if self.provider == "gemini":
return await self._call_gemini(messages, response_format, max_retries, initial_backoff, max_backoff, skip_validation, start_time, **kwargs)
call_params = {
"model": self.model,
@ -137,7 +151,6 @@ class LLMConfig:
if response_format is not None:
# Use JSON mode instead of strict parse for flexibility with optional fields
# This allows the LLM to omit optional fields without validation errors
import json
# Add schema to the system message
if hasattr(response_format, 'model_json_schema'):
@ -215,6 +228,128 @@ class LLMConfig:
raise last_exception
raise RuntimeError(f"LLM call failed after all retries with no exception captured")
async def _call_gemini(
self,
messages: List[Dict[str, str]],
response_format: Optional[Any],
max_retries: int,
initial_backoff: float,
max_backoff: float,
skip_validation: bool,
start_time: float,
**kwargs
) -> Any:
"""Handle Gemini-specific API calls using google-genai SDK."""
import json
# Convert OpenAI-style messages to Gemini format
# Gemini uses 'user' and 'model' roles, and system instructions are separate
system_instruction = None
gemini_contents = []
for msg in messages:
role = msg.get('role', 'user')
content = msg.get('content', '')
if role == 'system':
# Accumulate system messages as system instruction
if system_instruction:
system_instruction += "\n\n" + content
else:
system_instruction = content
elif role == 'assistant':
gemini_contents.append(genai_types.Content(
role="model",
parts=[genai_types.Part(text=content)]
))
else: # user or any other role
gemini_contents.append(genai_types.Content(
role="user",
parts=[genai_types.Part(text=content)]
))
# Add JSON schema instruction if response_format is provided
if response_format is not None and hasattr(response_format, 'model_json_schema'):
schema = response_format.model_json_schema()
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
if system_instruction:
system_instruction += schema_msg
else:
system_instruction = schema_msg
# Build generation config
config_kwargs = {}
if system_instruction:
config_kwargs['system_instruction'] = system_instruction
if 'temperature' in kwargs:
config_kwargs['temperature'] = kwargs['temperature']
if 'max_tokens' in kwargs:
config_kwargs['max_output_tokens'] = kwargs['max_tokens']
if response_format is not None:
config_kwargs['response_mime_type'] = 'application/json'
generation_config = genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
last_exception = None
for attempt in range(max_retries + 1):
try:
response = await self._gemini_client.aio.models.generate_content(
model=self.model,
contents=gemini_contents,
config=generation_config,
)
content = response.text
if response_format is not None:
# Parse the JSON response
json_data = json.loads(content)
# Return raw JSON if skip_validation is True, otherwise validate with Pydantic
if skip_validation:
result = json_data
else:
result = response_format.model_validate(json_data)
else:
result = content
# Log call details only if it takes more than 10 seconds
duration = time.time() - start_time
if duration > 10.0 and hasattr(response, 'usage_metadata') and response.usage_metadata:
usage = response.usage_metadata
logger.info(
f"slow llm call: model={self.provider}/{self.model}, "
f"input_tokens={usage.prompt_token_count}, output_tokens={usage.candidates_token_count}, "
f"time={duration:.3f}s"
)
return result
except genai_errors.APIError as e:
# Handle rate limits and server errors with retry
if e.code in (429, 503, 500):
last_exception = e
if attempt < max_retries:
backoff = min(initial_backoff * (2 ** attempt), max_backoff)
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
sleep_time = backoff + jitter
await asyncio.sleep(sleep_time)
else:
logger.error(f"Gemini API error after {max_retries + 1} attempts: {str(e)}")
raise
else:
logger.error(f"Gemini API error: {type(e).__name__}: {str(e)}")
raise
except Exception as e:
logger.error(f"Unexpected error during Gemini call: {type(e).__name__}: {str(e)}")
raise
if last_exception:
raise last_exception
raise RuntimeError(f"Gemini call failed after all retries with no exception captured")
@classmethod
def for_memory(cls) -> "LLMConfig":
"""Create configuration for memory operations from environment variables."""

1
hindsight-api/pyproject.toml
View file

@ -35,6 +35,7 @@ dependencies = [
"opentelemetry-instrumentation-fastapi>=0.41b0",
"opentelemetry-exporter-prometheus>=0.41b0",
"dateparser>=1.2.2",
"google-genai>=1.0.0",
]
[project.optional-dependencies]

81
hindsight-dev/benchmarks/longmemeval/longmemeval_benchmark.py
View file

@ -224,8 +224,10 @@ async def run_benchmark(
question_id: str = None,
only_failed: bool = False,
only_invalid: bool = False,
only_ingested: bool = False,
category: str = None,
max_concurrent_items: int = 1
max_concurrent_items: int = 1,
results_filename: str = "benchmark_results.json"
):
"""
Run the LongMemEval benchmark.
@ -241,8 +243,10 @@ async def run_benchmark(
question_id: Optional question ID to filter (e.g., 'e47becba'). Useful with --skip-ingestion.
only_failed: If True, only run questions that were previously marked as incorrect (is_correct=False)
only_invalid: If True, only run questions that were previously marked as invalid (is_invalid=True)
only_ingested: If True, only run questions whose memory bank already exists (has been ingested)
category: Optional category to filter questions (e.g., 'single-session-user', 'multi-session', 'temporal-reasoning'). Mutually exclusive with max_instances and max_instances_per_category.
max_concurrent_items: Maximum number of instances to process in parallel (default: 1 for sequential)
results_filename: Filename for results (default: benchmark_results.json). Directory is fixed to results/.
"""
from rich.console import Console
console = Console()
@ -253,6 +257,24 @@ async def run_benchmark(
console.print("[red]Error: --max-instances, --max-questions-per-category, and --category are mutually exclusive[/red]")
return
# Validate --only-ingested can't be combined with other dataset filters
if only_ingested:
incompatible_flags = []
if only_failed:
incompatible_flags.append("--only-failed")
if only_invalid:
incompatible_flags.append("--only-invalid")
if category is not None:
incompatible_flags.append("--category")
if question_id is not None:
incompatible_flags.append("--question-id")
if max_instances_per_category is not None:
incompatible_flags.append("--max-instances-per-category")
if incompatible_flags:
console.print(f"[red]Error: --only-ingested cannot be combined with: {', '.join(incompatible_flags)}[/red]")
return
# Check dataset exists, download if needed
dataset_path = Path(__file__).parent / "datasets" / "longmemeval_s_cleaned.json"
if not dataset_path.exists():
@ -373,6 +395,40 @@ async def run_benchmark(
from benchmarks.common.benchmark_runner import create_memory_engine
memory = await create_memory_engine()
# Filter by only_ingested: only run items whose memory bank already exists
if only_ingested:
console.print("[cyan]Filtering to only items with existing memory banks...[/cyan]")
# Load all items if not already loaded
if original_dataset_items is None:
original_dataset_items = dataset.load(dataset_path, max_items=None)
items_to_check = filtered_items if filtered_items is not None else original_dataset_items
# Check which items have existing banks
ingested_items = []
pool = await memory._get_pool()
for item in items_to_check:
item_id = dataset.get_item_id(item)
agent_id = f"longmemeval_{item_id}"
# Check if bank has any memory units
async with pool.acquire() as conn:
result = await conn.fetchrow(
"SELECT COUNT(*) as count FROM memory_units WHERE bank_id = $1 LIMIT 1",
agent_id
)
if result['count'] > 0:
ingested_items.append(item)
filtered_items = ingested_items
console.print(f"[green]Found {len(filtered_items)} items with existing memory banks[/green]")
if not filtered_items:
console.print("[yellow]No items found with existing memory banks. Nothing to run.[/yellow]")
return
# Create benchmark runner
runner = BenchmarkRunner(
dataset=dataset,
@ -381,7 +437,7 @@ async def run_benchmark(
memory=memory
)
# If filtering by category, failed, invalid, or max_instances_per_category, we need to use a custom dataset that only returns those items
# If filtering by category, failed, invalid, only_ingested, or max_instances_per_category, we need to use a custom dataset that only returns those items
# We'll temporarily replace the dataset's load method
if filtered_items is not None:
original_load = dataset.load
@ -392,12 +448,12 @@ async def run_benchmark(
# Run benchmark
# Single-phase approach: each question gets its own isolated agent_id
# This ensures each question only has access to its own context
output_path = Path(__file__).parent / 'results' / 'benchmark_results.json'
output_path = Path(__file__).parent / 'results' / results_filename
# Create results directory if it doesn't exist
output_path.parent.mkdir(parents=True, exist_ok=True)
merge_with_existing = (filln or question_id is not None or only_failed or only_invalid or category is not None or max_instances_per_category is not None)
merge_with_existing = (filln or question_id is not None or only_failed or only_invalid or only_ingested or category is not None or max_instances_per_category is not None)
results = await runner.run(
dataset_path=dataset_path,
@ -406,7 +462,7 @@ async def run_benchmark(
max_questions_per_item=max_questions_per_instance,
thinking_budget=thinking_budget,
max_tokens=max_tokens,
skip_ingestion=skip_ingestion,
skip_ingestion=skip_ingestion or only_ingested, # Auto-skip ingestion when using --only-ingested
max_concurrent_questions=8,
eval_semaphore_size=8,
separate_ingestion_phase=False, # Process each question independently
@ -576,6 +632,11 @@ if __name__ == "__main__":
action="store_true",
help="Only run questions that were previously marked as invalid (is_invalid=True). Requires existing results file."
)
parser.add_argument(
"--only-ingested",
action="store_true",
help="Only run questions whose memory bank already exists (has been ingested). Automatically skips ingestion. Cannot be combined with --only-failed, --only-invalid, --category, --question-id, or --max-instances-per-category."
)
parser.add_argument(
"--category",
type=str,
@ -588,6 +649,12 @@ if __name__ == "__main__":
default=1,
help="Number of instances to process in parallel (default: 1 for sequential). Higher values speed up evaluation but use more memory."
)
parser.add_argument(
"--results-filename",
type=str,
default="benchmark_results.json",
help="Filename for results output (default: benchmark_results.json). Saved in results/ directory."
)
args = parser.parse_args()
@ -615,6 +682,8 @@ if __name__ == "__main__":
question_id=args.question_id,
only_failed=args.only_failed,
only_invalid=args.only_invalid,
only_ingested=args.only_ingested,
category=args.category,
max_concurrent_items=args.parallel
max_concurrent_items=args.parallel,
results_filename=args.results_filename
))

2
hindsight/pyproject.toml
View file

@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "hindsight-all"
version = "0.0.7"
version = "0.0.17"
description = "All-in-one package for Hindsight - Semantic memory system with personality-driven thinking"
readme = "README.md"
requires-python = ">=3.11"

4
scripts/release.sh
View file

@ -65,7 +65,7 @@ fi
print_info "Updating version in all components..."
# Update Python packages
PYTHON_PACKAGES=("hindsight-api" "hindsight-dev" "hindsight-dev/benchmarks")
PYTHON_PACKAGES=("hindsight-api" "hindsight-dev" "hindsight-dev/benchmarks" "hindsight")
for package in "${PYTHON_PACKAGES[@]}"; do
PYPROJECT_FILE="$package/pyproject.toml"
if [ -f "$PYPROJECT_FILE" ]; then
@ -148,7 +148,7 @@ git add -A
git commit -m "Release v$VERSION
- Update version to $VERSION in all components
- Python packages: hindsight-api, hindsight-dev, hindsight-dev/benchmarks
- Python packages: hindsight-api, hindsight-dev, hindsight-dev/benchmarks, hindsight-all
- Python client: hindsight-clients/python
- TypeScript client: hindsight-clients/typescript
- Rust CLI: hindsight-cli

75
uv.lock
View file

@ -1027,6 +1027,44 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/01/61/d4b89fec821f72385526e1b9d9a3a0385dda4a72b206d28049e2c7cd39b8/gitpython-3.1.45-py3-none-any.whl", hash = "sha256:8908cb2e02fb3b93b7eb0f2827125cb699869470432cc885f019b8fd0fccff77", size = 208168 },
]
[[package]]
name = "google-auth"
version = "2.43.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "cachetools" },
{ name = "pyasn1-modules" },
{ name = "rsa" },
]
sdist = { url = "https://files.pythonhosted.org/packages/ff/ef/66d14cf0e01b08d2d51ffc3c20410c4e134a1548fc246a6081eae585a4fe/google_auth-2.43.0.tar.gz", hash = "sha256:88228eee5fc21b62a1b5fe773ca15e67778cb07dc8363adcb4a8827b52d81483", size = 296359 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/6f/d1/385110a9ae86d91cc14c5282c61fe9f4dc41c0b9f7d423c6ad77038c4448/google_auth-2.43.0-py2.py3-none-any.whl", hash = "sha256:af628ba6fa493f75c7e9dbe9373d148ca9f4399b5ea29976519e0a3848eddd16", size = 223114 },
]
[package.optional-dependencies]
requests = [
{ name = "requests" },
]
[[package]]
name = "google-genai"
version = "1.53.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "anyio" },
{ name = "google-auth", extra = ["requests"] },
{ name = "httpx" },
{ name = "pydantic" },
{ name = "requests" },
{ name = "tenacity" },
{ name = "typing-extensions" },
{ name = "websockets" },
]
sdist = { url = "https://files.pythonhosted.org/packages/de/b3/36fbfde2e21e6d3bc67780b61da33632f495ab1be08076cf0a16af74098f/google_genai-1.53.0.tar.gz", hash = "sha256:938a26d22f3fd32c6eeeb4276ef204ef82884e63af9842ce3eac05ceb39cbd8d", size = 260102 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/40/f2/97fefdd1ad1f3428321bac819ae7a83ccc59f6439616054736b7819fa56c/google_genai-1.53.0-py3-none-any.whl", hash = "sha256:65a3f99e5c03c372d872cda7419f5940e723374bb12a2f3ffd5e3e56e8eb2094", size = 262015 },
]
[[package]]
name = "greenlet"
version = "3.2.4"
@ -1117,7 +1155,7 @@ wheels = [
[[package]]
name = "hindsight-all"
version = "0.0.7"
version = "0.0.17"
source = { editable = "hindsight" }
dependencies = [
{ name = "hindsight-api" },
@ -1149,6 +1187,7 @@ dependencies = [
{ name = "dateparser" },
{ name = "fastapi", extra = ["standard"] },
{ name = "fastmcp" },
{ name = "google-genai" },
{ name = "greenlet" },
{ name = "httpx" },
{ name = "langchain-text-splitters" },
@ -1201,6 +1240,7 @@ requires-dist = [
{ name = "fastapi", extras = ["standard"], specifier = ">=0.120.3" },
{ name = "fastmcp", specifier = ">=2.0.0" },
{ name = "filelock", marker = "extra == 'test'", specifier = ">=3.0.0" },
{ name = "google-genai", specifier = ">=1.0.0" },
{ name = "greenlet", specifier = ">=3.2.4" },
{ name = "httpx", specifier = ">=0.27.0" },
{ name = "langchain-text-splitters", specifier = ">=0.3.0" },
@ -2870,6 +2910,27 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/e5/4e/519c1bc1876625fe6b71e9a28287c43ec2f20f73c658b9ae1d485c0c206e/pyarrow-21.0.0-cp313-cp313t-win_amd64.whl", hash = "sha256:222c39e2c70113543982c6b34f3077962b44fca38c0bd9e68bb6781534425c10", size = 26371006 },
]
[[package]]
name = "pyasn1"
version = "0.6.1"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/ba/e9/01f1a64245b89f039897cb0130016d79f77d52669aae6ee7b159a6c4c018/pyasn1-0.6.1.tar.gz", hash = "sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034", size = 145322 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/c8/f1/d6a797abb14f6283c0ddff96bbdd46937f64122b8c925cab503dd37f8214/pyasn1-0.6.1-py3-none-any.whl", hash = "sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629", size = 83135 },
]
[[package]]
name = "pyasn1-modules"
version = "0.4.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pyasn1" },
]
sdist = { url = "https://files.pythonhosted.org/packages/e9/e6/78ebbb10a8c8e4b61a59249394a4a594c1a7af95593dc933a349c8d00964/pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6", size = 307892 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/47/8d/d529b5d697919ba8c11ad626e835d4039be708a35b0d22de83a269a6682c/pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a", size = 181259 },
]
[[package]]
name = "pycparser"
version = "2.23"
@ -3628,6 +3689,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ed/d2/4a73b18821fd4669762c855fd1f4e80ceb66fb72d71162d14da58444a763/rpds_py-0.28.0-pp311-pypy311_pp73-musllinux_1_2_x86_64.whl", hash = "sha256:5d0145edba8abd3db0ab22b5300c99dc152f5c9021fab861be0f0544dc3cbc5f", size = 552199 },
]
[[package]]
name = "rsa"
version = "4.9.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pyasn1" },
]
sdist = { url = "https://files.pythonhosted.org/packages/da/8a/22b7beea3ee0d44b1916c0c1cb0ee3af23b700b6da9f04991899d0c555d4/rsa-4.9.1.tar.gz", hash = "sha256:e7bdbfdb5497da4c07dfd35530e1a902659db6ff241e39d9953cad06ebd0ae75", size = 29034 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/64/8d/0133e4eb4beed9e425d9a98ed6e081a55d195481b7632472be1af08d2f6b/rsa-4.9.1-py3-none-any.whl", hash = "sha256:68635866661c6836b8d39430f97a996acbd61bfa49406748ea243539fe239762", size = 34696 },
]
[[package]]
name = "safetensors"
version = "0.6.2"