From 5ad2dfe03eeab96e9be87366ae6a1dbb962642c1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Fri, 5 Dec 2025 01:00:25 +0100 Subject: [PATCH] gemini support --- .../hindsight_api/engine/llm_wrapper.py | 145 +++++++++++++++++- hindsight-api/pyproject.toml | 1 + .../longmemeval/longmemeval_benchmark.py | 81 +++++++++- hindsight/pyproject.toml | 2 +- scripts/release.sh | 4 +- uv.lock | 75 ++++++++- 6 files changed, 293 insertions(+), 15 deletions(-) diff --git a/hindsight-api/hindsight_api/engine/llm_wrapper.py b/hindsight-api/hindsight_api/engine/llm_wrapper.py index c38265f7..980ddad3 100644 --- a/hindsight-api/hindsight_api/engine/llm_wrapper.py +++ b/hindsight-api/hindsight_api/engine/llm_wrapper.py @@ -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.""" diff --git a/hindsight-api/pyproject.toml b/hindsight-api/pyproject.toml index fb707bf0..5c1417d4 100644 --- a/hindsight-api/pyproject.toml +++ b/hindsight-api/pyproject.toml @@ -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] diff --git a/hindsight-dev/benchmarks/longmemeval/longmemeval_benchmark.py b/hindsight-dev/benchmarks/longmemeval/longmemeval_benchmark.py index f612d10a..c66df7f4 100644 --- a/hindsight-dev/benchmarks/longmemeval/longmemeval_benchmark.py +++ b/hindsight-dev/benchmarks/longmemeval/longmemeval_benchmark.py @@ -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 )) diff --git a/hindsight/pyproject.toml b/hindsight/pyproject.toml index 02a7684f..ddaf8dae 100644 --- a/hindsight/pyproject.toml +++ b/hindsight/pyproject.toml @@ -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" diff --git a/scripts/release.sh b/scripts/release.sh index e3cfef20..a162b93a 100755 --- a/scripts/release.sh +++ b/scripts/release.sh @@ -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 diff --git a/uv.lock b/uv.lock index b0eece75..f5a5160b 100644 --- a/uv.lock +++ b/uv.lock @@ -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"