gemcode 0.2.2__py3-none-any.whl

gemcode/kairos_daemon.py

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+import asyncio
+import copy
+import sys
+import uuid
+from dataclasses import dataclass
+from typing import Awaitable, Callable
+
+from google.adk.runners import Runner
+
+from gemcode.config import GemCodeConfig
+from gemcode.capability_routing import apply_capability_routing
+from gemcode.model_routing import pick_effective_model
+from gemcode.invoke import run_turn
+from gemcode.session_runtime import create_runner
+
+
+def _events_to_text(events) -> str:
+  parts: list[str] = []
+  for event in events:
+    if not getattr(event, "content", None) or not getattr(
+      event.content, "parts", None
+    ):
+      continue
+    for part in event.content.parts:
+      if getattr(part, "text", None) and getattr(event, "author", None) != "user":
+        parts.append(part.text)
+  return "".join(parts)
+
+
+@dataclass(frozen=True)
+class KairosJob:
+  job_id: str
+  prompt: str
+  priority: int
+  session_id: str
+
+
+class KairosDaemon:
+  """Kairos-like proactive scheduler (stdin -> priority queue -> job runners)."""
+
+  def __init__(
+    self,
+    *,
+    cfg: GemCodeConfig,
+    concurrency: int = 2,
+    default_priority: int = 0,
+    user_id: str = "local",
+    job_runner: Callable[[KairosJob], Awaitable[None]] | None = None,
+  ) -> None:
+    self.cfg = cfg
+    self.concurrency = max(1, int(concurrency))
+    self.default_priority = int(default_priority)
+    self.user_id = user_id
+
+    # Queue items are (sort_key, seq, KairosJob).
+    self._queue: asyncio.PriorityQueue[
+      tuple[int, int, KairosJob]
+    ] = asyncio.PriorityQueue()
+    self._seq = 0
+    self._sem = asyncio.Semaphore(self.concurrency)
+    self._stop_event = asyncio.Event()
+
+    self._job_runner = job_runner or self._default_job_runner
+
+  def enqueue_prompt(
+    self,
+    *,
+    prompt: str,
+    priority: int | None = None,
+    session_id: str,
+  ) -> str:
+    """Enqueue a new job into the priority queue and return job_id."""
+    job_id = f"job_{uuid.uuid4().hex[:10]}"
+    pr = self.default_priority if priority is None else int(priority)
+    self._seq += 1
+    job = KairosJob(
+      job_id=job_id,
+      prompt=prompt,
+      priority=pr,
+      session_id=session_id,
+    )
+    # Higher priority should run first => use negative sort key.
+    self._queue.put_nowait((-pr, self._seq, job))
+    return job_id
+
+  async def _default_job_runner(self, job: KairosJob) -> None:
+    runner: Runner | None = None
+    try:
+      # Route model/capabilities based on this job's prompt, without mutating
+      # the daemon's base config shared across jobs.
+      job_cfg = copy.deepcopy(self.cfg)
+      apply_capability_routing(job_cfg, job.prompt, context="prompt")
+      job_cfg.model = pick_effective_model(job_cfg, job.prompt)
+
+      # For the initial MVP, we inject Kairos tools via `_build_extra_tools_for_job()`;
+      # this keeps scheduling logic independent from tool declarations.
+      extra_tools = self._build_extra_tools_for_job(job)
+      runner = create_runner(job_cfg, extra_tools=extra_tools or None)
+      events = await run_turn(
+        runner,
+        user_id=self.user_id,
+        session_id=job.session_id,
+        prompt=job.prompt,
+        max_llm_calls=job_cfg.max_llm_calls,
+        cfg=job_cfg,
+      )
+      text = _events_to_text(events).strip()
+      if text:
+        print(f"\n[kairos {job.job_id}] {text}\n", flush=True)
+      else:
+        print(f"\n[kairos {job.job_id}] (no text output)\n", flush=True)
+    finally:
+      if runner is not None:
+        await runner.close()
+
+  def _build_extra_tools_for_job(self, job: KairosJob) -> list | None:
+    """Inject per-job tools for the model to call."""
+
+    async def kairos_sleep_ms(duration_ms: int) -> dict:
+      """Pause this job for `duration_ms` (does not block other jobs)."""
+      duration_ms = max(0, int(duration_ms))
+      await asyncio.sleep(duration_ms / 1000.0)
+      return {"slept_ms": duration_ms}
+
+    def kairos_enqueue_prompt(
+      prompt: str,
+      priority: int = 0,
+      session_id: str | None = None,
+    ) -> dict:
+      """Enqueue a new Kairos job from the model.
+
+      If `session_id` is not provided, it defaults to the current job's
+      session_id.
+      """
+      sid = job.session_id if session_id is None else str(session_id)
+      enqueued_id = self.enqueue_prompt(
+        prompt=prompt,
+        priority=priority,
+        session_id=sid,
+      )
+      return {"enqueued_job_id": enqueued_id}
+
+    return [kairos_sleep_ms, kairos_enqueue_prompt]
+
+  async def _run_job_with_semaphore(self, job: KairosJob) -> None:
+    async with self._sem:
+      await self._job_runner(job)
+
+  async def _run_job_and_release(self, job: KairosJob) -> None:
+    try:
+      await self._job_runner(job)
+    finally:
+      self._sem.release()
+
+  async def drain(self, *, max_jobs: int | None = None) -> int:
+    """Process jobs already queued, useful for unit tests."""
+    processed = 0
+    while not self._queue.empty() and (max_jobs is None or processed < max_jobs):
+      _, _, job = await self._queue.get()
+      await self._run_job_with_semaphore(job)
+      processed += 1
+    return processed
+
+  async def _stdin_loop(self, *, session_id: str) -> None:
+    """Read stdin lines and enqueue each as a new job."""
+    # Use a background thread so the asyncio loop stays responsive.
+    prompt_prefix = "kairos> "
+    while not self._stop_event.is_set():
+      try:
+        # Print prompt only in interactive terminals.
+        if sys.stdin.isatty():
+          print(prompt_prefix, end="", flush=True)
+        line = await asyncio.to_thread(sys.stdin.readline)
+      except Exception:
+        break
+      if not line:
+        # EOF.
+        break
+
+      s = line.strip()
+      if not s:
+        continue
+      if s.lower() in ("quit", "exit", "q"):
+        self._stop_event.set()
+        break
+
+      # MVP: one line => one job at default priority.
+      self.enqueue_prompt(
+        prompt=s,
+        priority=self.default_priority,
+        session_id=session_id,
+      )
+
+  async def run_forever(self, *, session_id: str) -> None:
+    """Start the scheduler and keep running until stdin EOF/quit."""
+
+    scheduler_task = asyncio.create_task(self._scheduler_loop())
+    stdin_task = asyncio.create_task(self._stdin_loop(session_id=session_id))
+
+    # Wait for either scheduler to stop (shouldn't happen) or stdin loop to end.
+    done, pending = await asyncio.wait(
+      {scheduler_task, stdin_task},
+      return_when=asyncio.FIRST_COMPLETED,
+    )
+    for p in pending:
+      p.cancel()
+
+  async def _scheduler_loop(self) -> None:
+    """Continuously dequeue jobs by priority and run them."""
+    while not self._stop_event.is_set():
+      try:
+        # Don't dequeue from the priority queue unless we can start work
+        # immediately. This preserves priority ordering for "next starts".
+        await self._sem.acquire()
+        _, _, job = await self._queue.get()
+      except asyncio.CancelledError:
+        break
+      asyncio.create_task(self._run_job_and_release(job))
+

gemcode/limits.py

@@ -0,0 +1,83 @@
+"""
+Pre-model limits (cf. Claude `calculateTokenWarningState` / blocking limit checks).
+
+Uses session state updated in `callbacks.make_after_model_callback`.
+"""
+
+from __future__ import annotations
+
+from gemcode.config import GemCodeConfig
+
+SESSION_TOTAL_TOKENS_KEY = "gemcode:session_total_tokens"
+TOKEN_BUDGET_STOP_KEY = "gemcode:bt_token_budget_stop"
+TERMINAL_REASON_KEY = "gemcode:terminal_reason"
+
+
+def make_before_model_limits_callback(cfg: GemCodeConfig):
+  """Block the next LLM call when cumulative session tokens exceed ceiling."""
+  if cfg.max_session_tokens is None:
+    return None
+
+  async def before_model(callback_context, llm_request):
+    st = callback_context.state
+    total = int(st.get(SESSION_TOTAL_TOKENS_KEY, 0) or 0)
+    if total >= cfg.max_session_tokens:
+      from google.adk.models.llm_response import LlmResponse
+      from google.genai import types
+
+      # Record a terminal reason for stopHooks-like taxonomy.
+      callback_state = callback_context.state
+      if not callback_state.get(TERMINAL_REASON_KEY):
+        callback_state[TERMINAL_REASON_KEY] = "session_token_limit"
+
+      return LlmResponse(
+          content=types.Content(
+              role="model",
+              parts=[
+                  types.Part(
+                      text=(
+                          f"GemCode: session token ceiling ({cfg.max_session_tokens}) reached "
+                          "(see GEMCODE_MAX_SESSION_TOKENS). Start a new session or raise the limit."
+                      )
+                  )
+              ],
+          ),
+          turn_complete=True,
+      )
+    return None
+
+  return before_model
+
+
+def make_before_model_token_budget_callback(cfg: GemCodeConfig):
+  """Short-circuit the next model call after token-budget stop flag."""
+  if cfg.token_budget is None:
+    return None
+
+  async def before_model(callback_context, llm_request):
+    st = callback_context.state
+    if not st.get(TOKEN_BUDGET_STOP_KEY, False):
+      return None
+
+    from google.adk.models.llm_response import LlmResponse
+    from google.genai import types
+
+    if not st.get(TERMINAL_REASON_KEY):
+      st[TERMINAL_REASON_KEY] = "token_budget_stop"
+
+    return LlmResponse(
+      content=types.Content(
+        role="model",
+        parts=[
+          types.Part(
+            text=(
+              f"GemCode: token budget ({cfg.token_budget}) exhausted for this turn. "
+              "Start a new request to continue."
+            )
+          )
+        ],
+      ),
+      turn_complete=True,
+    )
+
+  return before_model

gemcode/live_audio_engine.py

@@ -0,0 +1,124 @@
+"""
+Live audio engine (Gemini Live API via ADK).
+
+This wires GemCode's existing outer session + callbacks into ADK's
+`Runner.run_live()` path for real-time audio input/output.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+from typing import Optional
+
+from google.adk.agents.live_request_queue import LiveRequestQueue
+from google.adk.agents.run_config import RunConfig
+from google.genai import types
+
+from gemcode.config import GemCodeConfig
+from gemcode.session_runtime import create_runner
+
+
+def _mime_type_for_rate(rate: int) -> str:
+  # ADK/examples commonly use this mime type.
+  return f"audio/pcm;rate={rate}"
+
+
+def _record_mic_pcm_blocking(*, rate: int, seconds: int) -> bytes:
+  try:
+    import sounddevice as sd
+    import numpy as np
+  except ImportError as e:
+    raise RuntimeError(
+      "Mic capture requires `sounddevice` and `numpy`. Install them to use `gemcode live-audio`."
+    ) from e
+
+  frames = int(rate * seconds)
+  # mono int16
+  audio = sd.rec(frames, samplerate=rate, channels=1, dtype="int16")
+  sd.wait()
+  pcm = np.asarray(audio).astype("int16", copy=False)
+  return pcm.tobytes()
+
+
+async def run_live_audio(
+  cfg: GemCodeConfig,
+  *,
+  session_id: str,
+  user_id: str = "local",
+  seconds: int = 10,
+  input_rate: int = 24_000,
+  language_code: Optional[str] = None,
+) -> None:
+  """
+  Record microphone audio for `seconds` and send it to Gemini Live.
+
+  MVP behavior:
+  - sends the entire recorded buffer as a single audio blob
+  - prints any model text parts it returns (typically transcriptions)
+  """
+
+  runner = create_runner(cfg)
+  live_queue = LiveRequestQueue()
+
+  speech_config = None
+  if language_code:
+    speech_config = types.SpeechConfig(language_code=language_code)
+
+  run_config = RunConfig(
+    response_modalities=["AUDIO"],
+    speech_config=speech_config,
+    # Keep SDK defaults for STT/TTS transcription configs.
+  )
+
+  agen = runner.run_live(
+    user_id=user_id,
+    session_id=session_id,
+    live_request_queue=live_queue,
+    run_config=run_config,
+  )
+
+  printed_any = False
+
+  async def _consume_events() -> None:
+    nonlocal printed_any
+    try:
+      async for event in agen:
+        if not event.content or not event.content.parts:
+          continue
+        for part in event.content.parts:
+          part_text = getattr(part, "text", None)
+          # We only print model-authored text to avoid echoing user input.
+          if part_text and getattr(event, "author", None) and event.author != "user":
+            sys.stdout.write(part_text)
+            sys.stdout.flush()
+            printed_any = True
+    except Exception:
+      # Runner/live failures are expected to be surfaced as terminal errors
+      # in session state + audit logs; don't crash the CLI.
+      raise
+
+  consumer_task = asyncio.create_task(_consume_events())
+
+  # Send "user started speaking" signal.
+  live_queue.send_activity_start()
+
+  pcm_bytes = await asyncio.to_thread(
+    _record_mic_pcm_blocking, rate=input_rate, seconds=seconds
+  )
+  live_queue.send_realtime(
+    types.Blob(data=pcm_bytes, mime_type=_mime_type_for_rate(input_rate))
+  )
+
+  # Send "user finished speaking" signal and close the queue.
+  live_queue.send_activity_end()
+  live_queue.close()
+
+  # Wait for event stream to drain.
+  await consumer_task
+
+  if not printed_any:
+    print("\n[gemcode live-audio] No model text received (audio may have been silent).")
+
+  await runner.close()
+

gemcode/mcp_loader.py

@@ -0,0 +1,57 @@
+"""
+Optional MCP toolsets from `.gemcode/mcp.json`.
+
+Schema (example):
+{
+  "servers": [
+    {
+      "name": "docs",
+      "stdio": { "command": "npx", "args": ["-y", "@some/mcp-server"] }
+    }
+  ]
+}
+
+Requires: pip install gemcode[mcp]
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from gemcode.config import GemCodeConfig
+
+
+def load_mcp_toolsets(cfg: GemCodeConfig) -> list:
+  path = cfg.project_root / ".gemcode" / "mcp.json"
+  if not path.is_file():
+    return []
+  try:
+    data = json.loads(path.read_text(encoding="utf-8"))
+  except json.JSONDecodeError as e:
+    raise ValueError(f"Invalid mcp.json: {e}") from e
+
+  try:
+    from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
+    from mcp import StdioServerParameters
+  except ImportError as e:
+    raise ImportError("Install MCP extras: pip install gemcode[mcp]") from e
+
+  servers = data.get("servers") or []
+  toolsets: list[Any] = []
+  for s in servers:
+    stdio = s.get("stdio") or {}
+    cmd = stdio.get("command")
+    args = stdio.get("args") or []
+    if not cmd:
+      continue
+    params = StdioServerParameters(command=cmd, args=args)
+    prefix = s.get("name") or "mcp"
+    toolsets.append(
+      McpToolset(
+        connection_params=params,
+        tool_name_prefix=prefix,
+      )
+    )
+  return toolsets