Update Gemini dependency

DESIGN.md

@@ -11,7 +11,7 @@ The library is designed around a few core concepts:
 - **SamplingStrategy**: Encapsulates persona selection, context/document sampling, and instruction dispatch logic. Infers target usage counts from stopping callbacks and configures sampling for both context and persona pipelines.
 - **QualityGate**: Wraps the evaluator (ConversationJudge) and retry logic for the `auto_improve` workflow. Determines whether a generated conversation meets quality thresholds or needs regeneration.
 - **PersonaGenerator**: Analyzes documents to generate diverse user personas, enhancing dataset variety.
-- **LLMProvider**: An abstraction over different language model providers (Gemini, OpenAI compatible).
+- **LLMProvider**: An abstraction over different language model providers (Gemini, OpenAI compatible). The `google-genai` package dependency is upgraded (`>=1.50.0`) with built-in support for Gemini thought signatures in multi-turn conversations and automatic reasoning extraction (`reasoning_content`) from candidate thought parts.
 - **EmbeddingProvider**: Async-first text embeddings (`async def embed(texts) -> list[list[float]]`). API backends (`OpenAIEmbeddingProvider`, `GeminiEmbeddingProvider`) use each vendor’s async client and `SmartKeyPool`; `ProcessEmbeddingProvider` runs SentenceTransformer in a `ProcessPoolExecutor` so the asyncio loop is not blocked by local inference (install the `embeddings-local` extra for `sentence-transformers`). Use `EmbeddingProviderFactory.create({...})` in `afterimage/providers/embedding_providers.py`.
 - **DatasetStorage**: An abstraction for storing and loading generated conversations and documents. It supports JSONL and SQL backends.
 - **Callbacks**: These allow for customization of the generation process.
@@ -54,14 +54,15 @@ The code is organized into the following directories and files:
         It also holds provider-aware concurrency defaults.
     - `evaluator.py`: `ConversationJudge` and embedding defaults for auto-improve.
     - `key_management.py`: Smart API key management with rate limiting.
+    - `logging.py`: Central logging utilities (`silence_noisy_third_party_loggers` with default `ERROR` level to suppress `google_genai` API key warnings and HTTP noise).
     - `monitoring.py`: Monitoring system implementation.
     - `persona_generator.py`: **[NEW]** Logic for generating personas from documents.
     - `prompts.py`: Default prompts (instruction generation, respondent persona etc.).
     - `quality.py`: Quality checking logic.
     - `retrievers.py`: Context retrieval strategies for RAG (`ContextRetriever`, `RetrievalResult`, optional `*_context_with_metadata`, `QdrantRetriever`, `StaticContextRetriever`, composite retrievers).
     - `storage.py`: Storage backends (JSONL, SQL).
     - `types.py`: Data models using Pydantic.
-    - `simula/`: **OpenSimula (experimental)** — `OpenSimula` orchestrator, `taxonomy_builder` (optional `show_progress` + tqdm), `cli_logging` (`configure_example_console`, `silence_noisy_third_party_loggers`), `sampling`, `meta_prompt`, `critics`, `double_critic`, `evaluation`, `document_context`, and `tasks/` (single QA, MCQ, multiturn instruction callback).
+    - `simula/`: **OpenSimula (experimental)** — `OpenSimula` orchestrator, `taxonomy_builder` (optional `show_progress` + tqdm), `cli_logging` (`configure_example_console`, `silence_noisy_third_party_loggers` with default `ERROR` level to mute `google_genai` API key warnings and HTTP noise), `sampling`, `meta_prompt`, `critics`, `double_critic`, `evaluation`, `document_context`, and `tasks/` (single QA, MCQ, multiturn instruction callback).
     - `evaluation/`: The new evaluation framework.
         - `__init__.py`: Exposes evaluation classes.
         - `base.py`: Base classes for evaluators.

afterimage/__init__.py

@@ -27,6 +27,7 @@
     default_embedding_provider_config,  # noqa
 )
 from afterimage.key_management import SmartKeyPool  # noqa
+from afterimage.logging import silence_noisy_third_party_loggers  # noqa
 from afterimage.orchestrator import Orchestrator  # noqa
 from afterimage.quality_gate import QualityGate, QualityResult  # noqa
 from afterimage.sampling import SamplingStrategy  # noqa

afterimage/cli.py

@@ -22,6 +22,9 @@
 @click.version_option(package_name="afterimage")
 def main():
     """AfterImage -- synthetic conversation dataset generator."""
+    from .logging import silence_noisy_third_party_loggers
+
+    silence_noisy_third_party_loggers()
 
 
 @main.command("agent-trace")

afterimage/logging.py

@@ -0,0 +1,24 @@
+"""Central logging utilities for AfterImage."""
+
+from __future__ import annotations
+
+import logging
+
+_NOISY_LOGGERS = (
+    "httpx",
+    "httpcore",
+    "httpcore.connection",
+    "httpcore.http11",
+    "google_genai",
+    "google_genai._api_client",
+    "google_genai.client",
+    "google_genai.models",
+    "google.auth",
+    "google.auth.transport",
+)
+
+
+def silence_noisy_third_party_loggers(level: int = logging.ERROR) -> None:
+    """Turn down chatty HTTP and google-genai log lines (warnings/info)."""
+    for name in _NOISY_LOGGERS:
+        logging.getLogger(name).setLevel(level)

afterimage/providers/llm_providers.py

@@ -12,6 +12,7 @@
 
 from ..common import default_safety_settings
 from ..key_management import SmartKeyPool
+from ..logging import silence_noisy_third_party_loggers
 from ..types import MODEL_PROVIDER_NAMES, ConversationEntry, ModelProviderName
 
 T = TypeVar("T", bound=BaseModel)
@@ -131,6 +132,45 @@ def _gemini_retry_delay(
     return random.uniform(base * 0.5, base * 1.5)
 
 
+def _extract_gemini_text_and_reasoning(response: Any) -> tuple[str, str | None]:
+    """Safely extract main output text and reasoning content from a Gemini response."""
+    try:
+        candidates = getattr(response, "candidates", None)
+        if not candidates:
+            return getattr(response, "text", "") or "", None
+
+        candidate = candidates[0]
+        content = getattr(candidate, "content", None)
+        if not content:
+            return getattr(response, "text", "") or "", None
+
+        parts = getattr(content, "parts", None)
+        if not parts:
+            return getattr(response, "text", "") or "", None
+
+        text_pieces: list[str] = []
+        thought_pieces: list[str] = []
+
+        for part in parts:
+            is_thought = getattr(part, "thought", False)
+            part_text = getattr(part, "text", None)
+            if part_text:
+                if is_thought:
+                    thought_pieces.append(part_text)
+                else:
+                    text_pieces.append(part_text)
+
+        main_text = (
+            "".join(text_pieces)
+            if text_pieces
+            else (getattr(response, "text", "") or "")
+        )
+        reasoning = "\n".join(thought_pieces).strip() if thought_pieces else None
+        return main_text, reasoning
+    except Exception:
+        return getattr(response, "text", "") or "", None
+
+
 class GeminiChatSession(ChatSession):
     """Gemini chat session implementation."""
 
@@ -172,14 +212,16 @@ def send_message(
                     raise
                 time.sleep(self._retry_delay(attempt))
 
+        text, reasoning = _extract_gemini_text_and_reasoning(response)
         return LLMResponse(
-            text=response.text,
+            text=text,
             finish_reason=str(response.candidates[0].finish_reason),
             prompt_token_count=response.usage_metadata.prompt_token_count,
             completion_token_count=response.usage_metadata.candidates_token_count,
             total_token_count=response.usage_metadata.total_token_count,
             model_name=self.model_name,
             raw_response=response,
+            reasoning_content=reasoning,
         )
 
     def close(self) -> None:
@@ -230,16 +272,18 @@ async def asend_message(
                     raise
                 await asyncio.sleep(self._retry_delay(attempt))
 
+        text, reasoning = _extract_gemini_text_and_reasoning(response)
         total_token_count = response.usage_metadata.total_token_count
         self.token_count = total_token_count
         return LLMResponse(
-            text=response.text,
+            text=text,
             finish_reason=str(response.candidates[0].finish_reason),
             prompt_token_count=response.usage_metadata.prompt_token_count,
             completion_token_count=response.usage_metadata.candidates_token_count,
             total_token_count=total_token_count,
             model_name=self.model_name,
             raw_response=response,
+            reasoning_content=reasoning,
         )
 
     async def aclose(self) -> None:
@@ -496,6 +540,7 @@ def __init__(
         self.retry_initial_delay = retry_initial_delay
         self.retry_max_delay = retry_max_delay
         self.kwargs = kwargs
+        silence_noisy_third_party_loggers()
 
     def _retry_delay(self, attempt: int) -> float:
         return _gemini_retry_delay(
@@ -535,14 +580,16 @@ def generate_content(
                     config=generation_config,
                 )
 
+                text, reasoning = _extract_gemini_text_and_reasoning(response)
                 return LLMResponse(
-                    text=response.text,
+                    text=text,
                     prompt_token_count=response.usage_metadata.prompt_token_count,
                     completion_token_count=response.usage_metadata.candidates_token_count,
                     total_token_count=response.usage_metadata.total_token_count,
                     finish_reason=str(response.candidates[0].finish_reason),
                     model_name=self.model_name,
                     raw_response=response,
+                    reasoning_content=reasoning,
                 )
 
             except Exception as exc:
@@ -584,14 +631,16 @@ async def agenerate_content(
                     config=generation_config,
                 )
 
+                text, reasoning = _extract_gemini_text_and_reasoning(response)
                 return LLMResponse(
-                    text=response.text,
+                    text=text,
                     prompt_token_count=response.usage_metadata.prompt_token_count,
                     completion_token_count=response.usage_metadata.candidates_token_count,
                     total_token_count=response.usage_metadata.total_token_count,
                     finish_reason=str(response.candidates[0].finish_reason),
                     model_name=self.model_name,
                     raw_response=response,
+                    reasoning_content=reasoning,
                 )
 
             except Exception as exc:
@@ -630,17 +679,22 @@ def generate_structured(
                     config=generation_config,
                 )
 
+                text, reasoning = _extract_gemini_text_and_reasoning(response)
+                parsed_val = (
+                    response.parsed
+                    if hasattr(response, "parsed") and response.parsed is not None
+                    else schema.model_validate_json(text)
+                )
                 return StructuredLLMResponse(
-                    text=response.text or "",
-                    parsed=response.parsed
-                    if hasattr(response, "parsed")
-                    else schema.model_validate_json(response.text),
+                    text=text,
+                    parsed=parsed_val,
                     prompt_token_count=response.usage_metadata.prompt_token_count,
                     completion_token_count=response.usage_metadata.candidates_token_count,
                     total_token_count=response.usage_metadata.total_token_count,
                     finish_reason=str(response.candidates[0].finish_reason),
                     model_name=self.model_name,
                     raw_response=response,
+                    reasoning_content=reasoning,
                 )
 
             except Exception as exc:
@@ -678,17 +732,22 @@ async def agenerate_structured(
                     contents=prompt,
                     config=generation_config,
                 )
+                text, reasoning = _extract_gemini_text_and_reasoning(response)
+                parsed_val = (
+                    response.parsed
+                    if hasattr(response, "parsed") and response.parsed is not None
+                    else schema.model_validate_json(text)
+                )
                 return StructuredLLMResponse(
-                    text=response.text or "",
-                    parsed=response.parsed
-                    if hasattr(response, "parsed")
-                    else schema.model_validate_json(response.text),
+                    text=text,
+                    parsed=parsed_val,
                     prompt_token_count=response.usage_metadata.prompt_token_count,
                     completion_token_count=response.usage_metadata.candidates_token_count,
                     total_token_count=response.usage_metadata.total_token_count,
                     finish_reason=str(response.candidates[0].finish_reason),
                     model_name=self.model_name,
                     raw_response=response,
+                    reasoning_content=reasoning,
                 )
 
             except Exception as exc:

afterimage/simula/cli_logging.py

@@ -4,22 +4,7 @@
 
 import logging
 
-_NOISY_LOGGERS = (
-    "httpx",
-    "httpcore",
-    "httpcore.connection",
-    "httpcore.http11",
-    "google_genai",
-    "google_genai.models",
-    "google.auth",
-    "google.auth.transport",
-)
-
-
-def silence_noisy_third_party_loggers(level: int = logging.WARNING) -> None:
-    """Turn down chatty HTTP and google-genai log lines during OpenSimula runs."""
-    for name in _NOISY_LOGGERS:
-        logging.getLogger(name).setLevel(level)
+from ..logging import silence_noisy_third_party_loggers
 
 
 def configure_example_console(

examples/agent_trace_generator_demo.py

@@ -98,7 +98,7 @@ async def main():
         teacher_model="gemini-3.5-flash-lite",
         judge_model="gemini-3.6-flash",
         observation_mode="llm",  # Preferred production mode (ESAT paper LLM observation synthesis).
-        task_synthesis_mode="simula",
+        task_synthesis_mode="grid",
         context_generator=VirtualUserContextGenerator(seed=42),
     )
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "afterimage"
-version = "0.18.0"
+version = "0.18.1"
 keywords = ["afterimage", "dataset-generation", "synthetic-dataset", "tool-calling", "structured-output", "preference-data", "persona-driven"]
 authors = [
     { name = "Yusuf Sarıgöz", email = "yusuf@altai.dev"}
@@ -26,7 +26,7 @@ dependencies = [
     "click>=8.0.0",
     "faker>=20.0.0",
     "filelock",
-    "google-genai==1.38.0",
+    "google-genai>=1.50.0",
     "grpcio==1.67.1",
     "huggingface-hub>=0.22.0",
     "langdetect>=1.0.9",