devtorch-core 3.0.1__py3-none-any.whl

devtorch_core/wrapper/gemini.py

@@ -0,0 +1,422 @@
+"""
+devtorch_core.wrapper.gemini
+=============================
+Wrapper around ``google.generativeai`` that injects RACP context and
+captures LLM reasoning to .GCC/.
+
+Because the Gemini API does not expose a separate ``system`` parameter in the
+same way Anthropic/OpenAI do, the RACP prefix is prepended to the text of
+the first user message part instead.
+
+Usage
+-----
+::
+
+    from devtorch_core.wrapper.gemini import DevTorchGemini
+
+    client = DevTorchGemini(model_name="gemini-1.5-pro")
+    response = client.generate_content(
+        [{"role": "user", "parts": [{"text": "Hello"}]}]
+    )
+"""
+from __future__ import annotations
+
+import logging
+import uuid
+from typing import Any, Iterator
+
+from .base import CaptureOrchestrator, RACPInjector, _reasoning_text, find_gcc_repo, is_disabled
+
+logger = logging.getLogger("devtorch.wrapper.gemini")
+
+# ---------------------------------------------------------------------------
+# Optional import
+# ---------------------------------------------------------------------------
+
+try:
+    import google.generativeai as _genai  # type: ignore[import]
+
+    HAS_GEMINI = True
+except ImportError:
+    HAS_GEMINI = False
+    _genai = None  # type: ignore[assignment]
+
+
+# ---------------------------------------------------------------------------
+# Wrapper
+# ---------------------------------------------------------------------------
+
+
+class DevTorchGemini:
+    """
+    Wrapper around ``google.generativeai.GenerativeModel``.
+
+    Provides a ``generate_content`` method that injects RACP context into the
+    first user message and captures the response to .GCC/.
+    """
+
+    def __init__(
+        self,
+        *args: Any,
+        model_name: str = "gemini-1.5-pro",
+        gcc_repo: Any = None,
+        **kwargs: Any,
+    ) -> None:
+        if not HAS_GEMINI:
+            raise ImportError(
+                "google-generativeai package not installed. "
+                "pip install devtorch[wrapper]"
+            )
+        self._model = _genai.GenerativeModel(model_name, **kwargs)  # type: ignore[union-attr]
+        self._model_name = model_name
+        self._gcc = gcc_repo or find_gcc_repo()
+        self._injector = RACPInjector(self._gcc) if self._gcc else None
+        self._capture = CaptureOrchestrator(self._gcc) if self._gcc else None
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def generate_content(
+        self,
+        messages: list,
+        stream: bool = False,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Inject RACP prefix into the first user message part, call the
+        GenerativeModel, and capture the response.
+
+        If ``stream=True`` yields chunks while accumulating the full text
+        for capture.
+
+        DevTorch errors never propagate to the caller.
+        """
+        session_id = str(uuid.uuid4())[:8]
+        call_id = str(uuid.uuid4())
+        prompt_text = _extract_prompt_text(messages)
+        injected_messages = self._inject(list(messages), prompt_text, session_id, call_id, self._model_name)
+
+        if stream:
+            return self._stream_generate(
+                injected_messages, session_id,
+                call_id=call_id,
+                prompt_text=prompt_text,
+                model_name=self._model_name,
+                **kwargs,
+            )
+
+        # Non-streaming
+        try:
+            response = self._model.generate_content(injected_messages, **kwargs)
+        except Exception:
+            self._record_llm_outcome(
+                prompt_text=prompt_text,
+                response_text="",
+                thinking_blocks=[],
+                session_id=session_id,
+                call_id=call_id,
+                model_name=self._model_name,
+                outcome="failure",
+            )
+            raise
+
+        text, thinking_blocks = self._capture_response(response, session_id)
+        self._record_llm_outcome(
+            prompt_text=prompt_text,
+            response_text=text,
+            thinking_blocks=thinking_blocks,
+            session_id=session_id,
+            call_id=call_id,
+            model_name=self._model_name,
+            outcome="success",
+        )
+
+        # Strip RACP blocks from text parts
+        try:
+            response = self._strip_response(response)
+        except Exception as exc:
+            logger.warning("devtorch: strip_response failed — %s", exc)
+
+        return response
+
+    # ------------------------------------------------------------------
+    # Streaming helper
+    # ------------------------------------------------------------------
+
+    def _stream_generate(
+        self,
+        messages: list,
+        session_id: str,
+        call_id: str = "",
+        prompt_text: str = "",
+        model_name: str = "",
+        **kwargs: Any,
+    ) -> "_GeminiStreamingWrapper":
+        raw_stream = self._model.generate_content(messages, stream=True, **kwargs)
+        return _GeminiStreamingWrapper(
+            raw_stream, self, session_id,
+            call_id=call_id,
+            prompt_text=prompt_text,
+            model_name=model_name,
+        )
+
+    # ------------------------------------------------------------------
+    # Injection
+    # ------------------------------------------------------------------
+
+    def _inject(
+        self,
+        messages: list,
+        prompt_text: str = "",
+        session_id: str = "",
+        call_id: str = "",
+        model_name: str = "",
+    ) -> list:
+        """Prepend RACP prefix to the first user message's first text part."""
+        if not self._injector or is_disabled():
+            return messages
+
+        try:
+            prefix, _ = self._injector.build_system_prefix_with_provenance(
+                prompt_text=prompt_text or None,
+                session_id=session_id,
+                call_id=call_id,
+                model_name=model_name,
+            )
+            if not prefix:
+                return messages
+            if (
+                messages
+                and isinstance(messages[0], dict)
+                and "parts" in messages[0]
+                and messages[0]["parts"]
+            ):
+                first_part = messages[0]["parts"][0]
+                if isinstance(first_part, dict) and "text" in first_part:
+                    messages[0]["parts"][0]["text"] = (
+                        prefix + "\n\n" + first_part["text"]
+                    )
+                elif isinstance(first_part, str):
+                    messages[0]["parts"][0] = prefix + "\n\n" + first_part
+        except Exception as exc:
+            logger.warning("devtorch: RACP injection failed — %s", exc)
+
+        return messages
+
+    # ------------------------------------------------------------------
+    # Capture
+    # ------------------------------------------------------------------
+
+    def _capture_response(self, response: Any, session_id: str) -> tuple[str, list]:
+        text = ""
+        thinking: list = []
+        if not self._capture or is_disabled():
+            return text, thinking
+        try:
+            text = _extract_text(response)
+            self._capture.capture(text, thinking, session_id)
+        except Exception as exc:
+            logger.warning("devtorch: _capture_response failed — %s", exc)
+        return text, thinking
+
+    def _record_llm_outcome(
+        self,
+        *,
+        prompt_text: str,
+        response_text: str,
+        thinking_blocks: list,
+        session_id: str,
+        call_id: str,
+        model_name: str,
+        outcome: str,
+    ) -> None:
+        if not self._capture or is_disabled():
+            return
+        try:
+            reasoning = _reasoning_text(thinking_blocks)
+            self._capture.record_learning_call(
+                call_type="llm",
+                reasoning=reasoning,
+                session_id=session_id,
+                call_id=call_id,
+                name=model_name,
+                input=prompt_text,
+                output=response_text,
+                outcome=outcome,
+            )
+        except Exception as exc:
+            logger.warning("devtorch: _record_llm_outcome failed — %s", exc)
+
+    # ------------------------------------------------------------------
+    # Strip
+    # ------------------------------------------------------------------
+
+    def _strip_response(self, response: Any) -> Any:
+        try:
+            from devtorch_core.parser.blocks import strip_racp_blocks
+            from devtorch_core.reasoning_plus.capture import strip_inline_thinking
+
+            # Gemini response.text is a property; try to mutate via parts
+            candidates = getattr(response, "candidates", None)
+            if candidates:
+                for candidate in candidates:
+                    content = getattr(candidate, "content", None)
+                    if content:
+                        parts = getattr(content, "parts", None)
+                        if parts:
+                            for part in parts:
+                                if hasattr(part, "text") and isinstance(part.text, str):
+                                    part.text = strip_inline_thinking(strip_racp_blocks(part.text))
+        except Exception as exc:
+            logger.warning("devtorch: strip_response failed — %s", exc)
+        return response
+
+    # Proxy everything else to the underlying model
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._model, name)
+
+
+# ---------------------------------------------------------------------------
+# Streaming wrapper
+# ---------------------------------------------------------------------------
+
+
+class _GeminiStreamingWrapper:
+    """
+    Thin wrapper around the raw Gemini streaming response.
+    """
+
+    def __init__(
+        self,
+        raw_stream: Any,
+        wrapper: "DevTorchGemini",
+        session_id: str,
+        call_id: str = "",
+        prompt_text: str = "",
+        model_name: str = "",
+    ) -> None:
+        self._raw = raw_stream
+        self._wrapper = wrapper
+        self._session_id = session_id
+        self._call_id = call_id
+        self._prompt_text = prompt_text
+        self._model_name = model_name
+        self._accumulated_text = ""
+        self._recorded = False
+
+    def __iter__(self) -> Iterator[Any]:
+        try:
+            for chunk in self._raw:
+                try:
+                    chunk_text = _extract_text(chunk)
+                    self._accumulated_text += chunk_text
+                except Exception:
+                    pass
+                yield chunk
+        except Exception:
+            self._on_close(outcome="failure")
+            raise
+        finally:
+            if not self._recorded:
+                self._on_close(outcome="success")
+
+    def __enter__(self) -> "_GeminiStreamingWrapper":
+        if hasattr(self._raw, "__enter__"):
+            self._raw.__enter__()
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        try:
+            if hasattr(self._raw, "__exit__"):
+                self._raw.__exit__(*args)
+        finally:
+            if not self._recorded:
+                self._on_close(outcome="success")
+
+    def _on_close(self, outcome: str = "success") -> None:
+        if self._recorded:
+            return
+        self._recorded = True
+        if self._accumulated_text:
+            try:
+                self._wrapper._capture_response(
+                    type("_R", (), {"text": self._accumulated_text})(),
+                    self._session_id,
+                )
+            except Exception as exc:
+                logger.warning("devtorch: stream capture failed — %s", exc)
+        try:
+            self._wrapper._record_llm_outcome(
+                prompt_text=self._prompt_text,
+                response_text=self._accumulated_text,
+                thinking_blocks=[],
+                session_id=self._session_id,
+                call_id=self._call_id,
+                model_name=self._model_name,
+                outcome=outcome,
+            )
+        except Exception as exc:
+            logger.warning("devtorch: stream learning outcome failed — %s", exc)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._raw, name)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _extract_text(response: Any) -> str:
+    """Pull plain text from a Gemini response object."""
+    # response.text is the simplest path
+    if hasattr(response, "text") and isinstance(response.text, str):
+        return response.text
+    # Fallback: iterate candidates → content → parts
+    text = ""
+    try:
+        candidates = getattr(response, "candidates", None) or []
+        for candidate in candidates:
+            content = getattr(candidate, "content", None)
+            if content:
+                parts = getattr(content, "parts", None) or []
+                for part in parts:
+                    if hasattr(part, "text") and isinstance(part.text, str):
+                        text += part.text
+    except Exception:
+        pass
+    return text
+
+
+def _extract_prompt_text(messages: list) -> str:
+    """Best-effort concatenation of user-visible message text for smart context."""
+    parts: list[str] = []
+    for m in messages or []:
+        if isinstance(m, dict):
+            role = m.get("role", "")
+            content = m.get("content", "")
+            parts_text = ""
+            if isinstance(content, str):
+                parts_text = content
+            elif isinstance(content, list):
+                for item in content:
+                    if isinstance(item, dict) and item.get("type") == "text":
+                        parts_text += item.get("text", "")
+                    elif isinstance(item, str):
+                        parts_text += item
+            elif isinstance(content, dict):
+                # Gemini shape: {"parts": [...]}
+                for part in content.get("parts", []):
+                    if isinstance(part, dict) and "text" in part:
+                        parts_text += part["text"]
+                    elif isinstance(part, str):
+                        parts_text += part
+            if role in ("user", "system") and parts_text:
+                parts.append(parts_text)
+        elif hasattr(m, "role") and hasattr(m, "content"):
+            role = getattr(m, "role", "")
+            content = getattr(m, "content", "")
+            if role in ("user", "system") and content:
+                parts.append(str(content))
+    return "\n".join(parts)[:4000]