diff --git a/.env.example b/.env.example index 23f6b13a..808480a8 100644 --- a/.env.example +++ b/.env.example @@ -21,11 +21,13 @@ CLIMATECLAW_VLLM_API_KEY="YOUR_VLLM_API_KEY" CLIMATECLAW_CODE_SERVER_PORT=8051 CLIMATECLAW_RAG_SERVER_PORT=8050 CLIMATECLAW_WEB_SEARCH_SERVER_PORT=8052 +CLIMATECLAW_PLUGIN_CODE_SEARCH_SERVER_PORT=8053 -CLIMATECLAW_AVAILABLE_MCP_SERVERS=rag-server,code-server,web-search-server # The list of servers to connect via backends client +CLIMATECLAW_AVAILABLE_MCP_SERVERS=rag-server,code-server,web-search-server,plugin-code-search-server # The list of servers to connect via backends client CLIMATECLAW_RAG_SERVER_URL=http://rag-server:8050 # IMPORTANT: Server URLs must be given as 'CLIMATECLAW_{server_name}_SERVER_URL' for successful parsing CLIMATECLAW_CODE_SERVER_URL=http://code-server:8051 CLIMATECLAW_WEB_SEARCH_SERVER_URL=http://web-search-server:8052 +CLIMATECLAW_PLUGIN_CODE_SEARCH_SERVER_URL=http://plugin-code-search-server:8053 CLIMATECLAW_MCP_REQUEST_TIMEOUT_SEC=600 # MONGODB SETTINGS diff --git a/.github/workflows/ci_job.yml b/.github/workflows/ci_job.yml index d1c95215..6248a0d4 100644 --- a/.github/workflows/ci_job.yml +++ b/.github/workflows/ci_job.yml @@ -108,6 +108,7 @@ jobs: CLIMATECLAW_MONGODB_PASSWORD=ci-psswrd CLIMATECLAW_MONGODB_DATABASE_NAME=ci-db CLIMATECLAW_OPENAI_API_KEY=${{ secrets.OPENAI_API_KEY }} + CLIMATECLAW_GITLAB_ACCESS_TOKEN=${{ secrets.GITLAB_ACCESS_TOKEN }} EOF - name: Build production base image diff --git a/.gitignore b/.gitignore index 6e289194..8e26996e 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,7 @@ wheels/ .coverage .env +scripts/evaluation/*.png # Virtual environments .venv diff --git a/README.md b/README.md index 937a6d2f..22a4c76e 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,8 @@ # ClimateClaw + ClimateClaw is a Python service for building AI-assisted climate-data workflows. It provides the API, conversation handling, model prompting, persistent thread storage, and tool orchestration needed to support interactive work with climate data. -The project integrates LiteLLM-native prompting, MongoDB-backed conversation state, and MCP-based tool execution for retrieval, code execution, and domain-specific automation. +The project integrates LiteLLM-native prompting, MongoDB-backed conversation states, and MCP-based tool execution for code execution, web/documentation search, code retrieval, and domain-specific automation. ## Highlights - FastAPI app with strict auth parity to the production Rust service (`/api/chatbot/*`) @@ -29,7 +30,6 @@ Services that start: - `climateclaw`: FastAPI app (debugpy toggle via `DEBUG=true` for remote debugging session) - `code-server`: MCP server running the sandboxed Jupyter kernel and exposing `code_interpreter` - `web-search-server`: MCP server doing web search via OpenAI API and exposing `web_search` -- `rag-server`: MCP server exposing `get_context_from_resources` - `litellm`: LiteLLM proxy that reads `litellm_config.yaml` Bind mounts expose `/work`, logs, threads, and shared `cache` to other Freva services. @@ -103,6 +103,7 @@ Create `.env` (used by FastAPI, Docker, and MCP servers). See `.env.example` for | `litellm_config.yaml` | Source of truth for model catalog (consumed by `available_chatbots()`) | Generated artifacts that persist across runs: + - `threads/` (JSONL transcript per thread id) - `cache/{user_id}/{thread_id}` (LLM-created files, plots, etc.) - `logs/` (when mounted in Docker) @@ -146,6 +147,7 @@ Generated artifacts that persist across runs: - **Code interpreter** (`src/climateclaw/tools/code/server.py`): spins up per-session Jupyter kernels, sanitizes input, enforces configurable timeouts, and injects Freva config via environment variables. Outputs include stdout/stderr, display data, and structured errors. - **Web search server** (`src/climateclaw/tools/web_search/server.py`): calls OpenAI Web Search (`gpt-4.1`) constrained to ICON model + DKRZ/HPC docs. Honors request cancellation. - **RAG server** (`src/climateclaw/tools/rag/server.py`): indexes documentation with custom loaders + splitters, stores embeddings in MongoDB (`embeddings`), and surfaces a single tool `get_context_from_resources`. LiteLLM requests embed queries through the same proxy (`CLIMATECLAW_LITE_LLM_ADDRESS`). +- **Plugin code search server** (`src/climateclaw/tools/plugin_code_search/server.py`): retrieves code from GitLab repositories used as Freva plugins for climate data analysis (included projects: *Coming Decade*, *ClimXtreme*, *RegiKlim*). Returns scraped code, doc and/or config files as query-relevant context. - **Header gate** (`src/climateclaw/tools/header_gate.py`): wraps each MCP ASGI app so critical headers become ContextVars and requests fail fast when missing/invalid (e.g., missing Mongo URI yields SSE-friendly JSON-RPC errors). - **Manager** (`src/climateclaw/services/mcp/mcp_manager.py`): caches clients, discovers tool schemas, exports OpenAI function definitions, and pins MCP session ids to thread ids for deterministic tool contexts. diff --git a/docker/docker-compose.dev.yml b/docker/docker-compose.dev.yml index dfd095fe..248c4db3 100644 --- a/docker/docker-compose.dev.yml +++ b/docker/docker-compose.dev.yml @@ -27,7 +27,11 @@ services: networks: - climateclaw healthcheck: - test: ["CMD-SHELL", "python -c \"import urllib.request; urllib.request.urlopen('http://127.0.0.1:${CLIMATECLAW_BACKEND_PORT:-8502}/healthz')\""] + test: + [ + "CMD-SHELL", + 'python -c "import urllib.request; urllib.request.urlopen(''http://127.0.0.1:${CLIMATECLAW_BACKEND_PORT:-8502}/healthz'')"', + ] interval: 30s timeout: 5s retries: 3 @@ -77,7 +81,11 @@ services: networks: - climateclaw healthcheck: - test: ["CMD-SHELL", "python -c \"import urllib.request; urllib.request.urlopen('http://127.0.0.1:8051/healthz')\""] + test: + [ + "CMD-SHELL", + 'python -c "import urllib.request; urllib.request.urlopen(''http://127.0.0.1:8051/healthz'')"', + ] interval: 30s timeout: 5s retries: 3 @@ -102,14 +110,48 @@ services: volumes: - ../logs/:/app/logs healthcheck: - test: ["CMD-SHELL", "python -c \"import urllib.request; urllib.request.urlopen('http://127.0.0.1:8052/healthz')\""] + test: + [ + "CMD-SHELL", + 'python -c "import urllib.request; urllib.request.urlopen(''http://127.0.0.1:8052/healthz'')"', + ] + interval: 30s + timeout: 5s + retries: 3 + start_period: 10s + + plugin-code-search-server: + image: plugin-code-search-server-${CLIMATECLAW_INSTANCE_NAME} + build: + context: .. + dockerfile: docker/plugin-code-search-server/Dockerfile + env_file: ../.env + hostname: plugin-code-search-server-${CLIMATECLAW_INSTANCE_NAME} + environment: + CLIMATECLAW_DEV: ${CLIMATECLAW_DEV} + CLIMATECLAW_DEBUG: ${CLIMATECLAW_DEBUG:-0} + CLIMATECLAW_MCP_HOST: 0.0.0.0 + CLIMATECLAW_MCP_PORT: 8053 + CLIMATECLAW_GITLAB_ACCESS_TOKEN: ${CLIMATECLAW_GITLAB_ACCESS_TOKEN:-} + ports: + - 8053:8053 + networks: + - climateclaw + volumes: + - ../logs/:/app/logs + healthcheck: + test: + [ + "CMD-SHELL", + 'python -c "import urllib.request; urllib.request.urlopen(''http://127.0.0.1:8053/healthz'')"', + ] interval: 30s timeout: 5s retries: 3 start_period: 10s litellm: - image: ghcr.io/berriai/litellm:main-latest + image: ghcr.io/berriai/litellm:v1.102.0 hostname: litellm-${CLIMATECLAW_INSTANCE_NAME} volumes: - ../litellm_config.dev.yaml:/app/config.yaml # Development config excludes the deployment-only vLLM model @@ -129,7 +171,11 @@ services: networks: - climateclaw healthcheck: - test: ["CMD-SHELL", "python -c \"import urllib.request; urllib.request.urlopen('http://localhost:4000/health/liveliness')\""] + test: + [ + "CMD-SHELL", + 'python -c "import urllib.request; urllib.request.urlopen(''http://localhost:4000/health/liveliness'')"', + ] interval: 30s timeout: 10s retries: 3 @@ -209,7 +255,7 @@ services: user: 0:0 hostname: volume-prep restart: "no" - command: find /settings -mindepth 1 -maxdepth 1 -exec rm -rf {} + + command: find /settings -mindepth 1 -maxdepth 1 -exec rm -rf {} + volumes: - settings-data:/settings:rw @@ -221,7 +267,14 @@ services: depends_on: volume-prep: condition: service_completed_successfully - command: ["clone", "--depth", "1", "https://github.com/freva-org/freva-service-config.git", "/settings"] + command: + [ + "clone", + "--depth", + "1", + "https://github.com/freva-org/freva-service-config.git", + "/settings", + ] volumes: - settings-data:/settings:rw @@ -234,11 +287,11 @@ services: git-prep: condition: service_completed_successfully command: > - sh -eu -c ' - ls -rlt /settings && - python /settings/dev-utils.py gen-certs --cert-dir /certs && - cp -a /settings/keycloak/import/. /import/ - ' + sh -eu -c ' + ls -rlt /settings && + python /settings/dev-utils.py gen-certs --cert-dir /certs && + cp -a /settings/keycloak/import/. /import/ + ' volumes: - settings-data:/settings:ro - certs-data:/certs:rw @@ -281,11 +334,11 @@ services: - "8080:8080" - "8443:8443" command: | - start-dev - --hostname-strict=false - --hostname=http://keycloak.localhost:8080 - --import-realm - -Dkeycloak.migration.strategy=OVERWRITE_EXISTING + start-dev + --hostname-strict=false + --hostname=http://keycloak.localhost:8080 + --import-realm + -Dkeycloak.migration.strategy=OVERWRITE_EXISTING keycloak-bootstrap: image: docker.io/alpine @@ -348,12 +401,10 @@ services: volumes: - certs-data:/certs:ro - networks: climateclaw: driver: bridge - volumes: settings-data: certs-data: diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index dff70efd..25fce08f 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -1,5 +1,4 @@ services: - climateclaw-base: # Build manually with: # docker compose -f docker/docker-compose.dev.yml --profile build-only build climateclaw-base @@ -26,7 +25,11 @@ services: networks: - climateclaw healthcheck: - test: ["CMD-SHELL", "python -c \"import urllib.request; urllib.request.urlopen('http://127.0.0.1:${CLIMATECLAW_BACKEND_PORT:-8502}/healthz')\""] + test: + [ + "CMD-SHELL", + 'python -c "import urllib.request; urllib.request.urlopen(''http://127.0.0.1:${CLIMATECLAW_BACKEND_PORT:-8502}/healthz'')"', + ] interval: 5s timeout: 3s retries: 20 @@ -53,7 +56,11 @@ services: networks: - climateclaw healthcheck: - test: ["CMD-SHELL", "python -c \"import urllib.request; urllib.request.urlopen('http://127.0.0.1:8051/healthz')\""] + test: + [ + "CMD-SHELL", + 'python -c "import urllib.request; urllib.request.urlopen(''http://127.0.0.1:8051/healthz'')"', + ] interval: 5s timeout: 3s retries: 20 @@ -74,7 +81,11 @@ services: volumes: - /container/da/climateclaw-links/${CLIMATECLAW_INSTANCE_NAME}/logs:/app/logs healthcheck: - test: ["CMD-SHELL", "python -c \"import urllib.request; urllib.request.urlopen('http://127.0.0.1:8052/healthz')\""] + test: + [ + "CMD-SHELL", + 'python -c "import urllib.request; urllib.request.urlopen(''http://127.0.0.1:8052/healthz'')"', + ] interval: 5s timeout: 3s retries: 20 @@ -102,7 +113,11 @@ services: networks: - climateclaw healthcheck: - test: ["CMD-SHELL", "python -c \"import urllib.request; urllib.request.urlopen('http://localhost:4000/health/liveliness')\""] + test: + [ + "CMD-SHELL", + 'python -c "import urllib.request; urllib.request.urlopen(''http://localhost:4000/health/liveliness'')"', + ] interval: 30s timeout: 10s retries: 3 diff --git a/docker/plugin-code-search-server/Dockerfile b/docker/plugin-code-search-server/Dockerfile new file mode 100644 index 00000000..ed600540 --- /dev/null +++ b/docker/plugin-code-search-server/Dockerfile @@ -0,0 +1,21 @@ +ARG VERSION=latest +FROM ghcr.io/freva-org/climateclaw-base:${VERSION} + +COPY docker/plugin-code-search-server/pyproject.toml /app/pyproject.toml + +# Export requirements from pyproject.toml into a file +RUN conda run -n env uv export --format requirements-txt --no-hashes -o /tmp/requirements.txt +# Install exactly those requirements into the env's interpreter +RUN conda run -n env uv pip install --no-cache-dir --system \ + --python "/opt/conda/envs/env/bin/python" \ + -r /tmp/requirements.txt + +EXPOSE ${CLIMATECLAW_MCP_PORT:-8053} + +COPY src /app/src + +# Launch server +CMD uvicorn climateclaw.tools.plugin_code_search.server:app \ + --host ${CLIMATECLAW_MCP_HOST:-0.0.0.0} \ + --port ${CLIMATECLAW_MCP_PORT:-8053} \ + --proxy-headers diff --git a/docker/plugin-code-search-server/pyproject.toml b/docker/plugin-code-search-server/pyproject.toml new file mode 100644 index 00000000..e1a54b35 --- /dev/null +++ b/docker/plugin-code-search-server/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "plugin-code-search-server" +version = "0.1.0" +description = "Add your description here" +readme = "README.md" +requires-python = ">=3.13" +dependencies = [ + "fastapi>=0.116.1", + "fastmcp>=2.12.4", + "httpx>=0.28.1", + "uvicorn[standard]>=0.35.0", + "openai", + "requests>=2.31.0", +] diff --git a/pyproject.toml b/pyproject.toml index 3435bf4a..4ec5dd3b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ dev = [ "pytest>=8.4.2", "pytest-asyncio>=1.2.0", "pytest-cov>=6.0.0", + "matplotlib>=3.11.1", ] test = [ "jupyter-client>=8.8.0", diff --git a/scripts/dev_benchmark_script.py b/scripts/dev_benchmark_script.py index 15575592..86413243 100644 --- a/scripts/dev_benchmark_script.py +++ b/scripts/dev_benchmark_script.py @@ -1,24 +1,23 @@ from __future__ import annotations +import asyncio +import json import os -import pathlib import sys +import time +from dataclasses import dataclass +from pathlib import Path -sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1])) +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) os.environ["CLIMATECLAW_DEV"] = "1" os.environ["CLIMATECLAW_LITE_LLM_ADDRESS"] = "http://localhost:4000" os.environ["CLIMATECLAW_RAG_SERVER_URL"] = "http://localhost:8050" os.environ["CLIMATECLAW_CODE_SERVER_URL"] = "http://localhost:8051" os.environ["CLIMATECLAW_WEB_SEARCH_SERVER_URL"] = "http://localhost:8052" +os.environ["CLIMATECLAW_PLUGIN_CODE_SEARCH_SERVER_URL"] = "http://localhost:8053" os.environ["CLIMATECLAW_MONGODB_HOST"] = "localhost" - -import asyncio -import json -import time -from dataclasses import dataclass - from climateclaw.core.logging_setup import configure_logging from climateclaw.core.prompting import get_entire_prompt from climateclaw.services.authentication.auth import Authenticator @@ -35,7 +34,7 @@ from climateclaw.services.streaming.stream_variants import from_sv_to_json """ -Headless dev/benchmark runner mirroring /chatbot/streamresponse behaviour. +Headless dev/benchmark runner mirroring /chatbot/streamresponse behaviour. Headless means that it runs without a web server, and without any user interaction. It is intended for development and benchmarking purposes. - Config below (no argparse). - Disk-only persistence (never Mongo). @@ -47,12 +46,12 @@ # CONFIG # ────────────────────────────────────────────────────────────────────────────── -MODEL = "gpt-4.1" +MODEL = "gpt-5.6-luna" # model to benchmark USER_ID = "janedoe" -PROMPT = "plot x=y" +PROMPT = "Hi, who are you?" -RUNS = 3 -CONCURRENCY = 3 # ← set to 1 for clean mode +RUNS = 1 # number of runs to perform +CONCURRENCY = 1 # ← set to 1 for clean mode PRINT_STREAM = False PRINT_PER_RUN_SUMMARY = True @@ -67,6 +66,7 @@ class RunResult: idx: int thread_id: str + answer: str duration_s: float chunks: int chars: int @@ -96,6 +96,7 @@ async def _run_once(idx: int, sem: asyncio.Semaphore) -> RunResult: chunk_count = 0 char_count = 0 status = "Done" + answer = "" try: async for variant in run_stream( @@ -109,6 +110,7 @@ async def _run_once(idx: int, sem: asyncio.Semaphore) -> RunResult: txt = getattr(variant, "text", "") or "" chunk_count += 1 char_count += len(txt) + answer += txt if PRINT_STREAM and getattr(variant, "variant", None) != "Assistant": print(json.dumps(from_sv_to_json(variant), ensure_ascii=False)) @@ -128,7 +130,15 @@ async def _run_once(idx: int, sem: asyncio.Semaphore) -> RunResult: f"chunks={chunk_count} chars={char_count} time={duration:.3f}s" ) - return RunResult(idx, thread_id, duration, chunk_count, char_count, status) + return RunResult( + idx, + thread_id, + answer, + duration, + chunk_count, + char_count, + status, + ) async def main() -> None: @@ -136,7 +146,7 @@ async def main() -> None: tasks = [asyncio.create_task(_run_once(i, sem)) for i in range(RUNS)] results = await asyncio.gather(*tasks) - if PRINT_FINAL_SUMMARY and results: + if results: ok = [r for r in results if r.status == "Done"] errs = [r for r in results if r.status != "Done"] avg = sum(r.duration_s for r in results) / len(results) @@ -146,17 +156,22 @@ async def main() -> None: total_chunks = sum(r.chunks for r in results) total_chars = sum(r.chars for r in results) - print("\n=== Summary ===") - print(f"model={MODEL} runs={RUNS} concurrency={CONCURRENCY}") - print(f"success={len(ok)} errors={len(errs)}") - print( - f"avg_time={avg:.3f}s p50_time={p50:.3f}s fastest={fastest.duration_s:.3f}s slowest={slowest.duration_s:.3f}s" - ) - print(f"total_chunks={total_chunks} total_chars={total_chars}") - if errs: - print("errors:") - for r in errs[:10]: - print(f" run={r.idx} thread={r.thread_id} status={r.status}") + if PRINT_FINAL_SUMMARY: + print("\n=== Summary ===") + print(f"model={MODEL} runs={RUNS} concurrency={CONCURRENCY}") + print(f"success={len(ok)} errors={len(errs)}") + print( + f"avg_time={avg:.3f}s p50_time={p50:.3f}s fastest={fastest.duration_s:.3f}s slowest={slowest.duration_s:.3f}s" + ) + print(f"total_chunks={total_chunks} total_chars={total_chars}") + print( + f"Assistant answer sample (from run {fastest.idx}):\n\n{fastest.answer}" + ) + + if errs: + print("errors:") + for r in errs[:10]: + print(f" run={r.idx} thread={r.thread_id} status={r.status}") if __name__ == "__main__": diff --git a/scripts/dev_chat.py b/scripts/dev_chat.py index 0aff8d3c..f48a1dd2 100644 --- a/scripts/dev_chat.py +++ b/scripts/dev_chat.py @@ -1,16 +1,17 @@ from __future__ import annotations import os -import pathlib import sys +from pathlib import Path -sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1])) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) os.environ["CLIMATECLAW_DEV"] = "1" os.environ["CLIMATECLAW_LITE_LLM_ADDRESS"] = "http://localhost:4000" os.environ["CLIMATECLAW_RAG_SERVER_URL"] = "http://localhost:8050" os.environ["CLIMATECLAW_CODE_SERVER_URL"] = "http://localhost:8051" os.environ["CLIMATECLAW_WEB_SEARCH_SERVER_URL"] = "http://localhost:8052" +os.environ["CLIMATECLAW_PLUGIN_CODE_SEARCH_SERVER_URL"] = "http://localhost:8053" os.environ["CLIMATECLAW_MONGODB_HOST"] = "localhost" @@ -69,7 +70,7 @@ # CONFIG # ────────────────────────────────────────────────────────────────────────────── -MODEL = "gpt-4.1" +MODEL = "gpt-5.6-luna" # model to use USER_ID = "janedoe" PRINT_DEBUG = False # Print non-Assistant stream variants (ServerHint, etc.) diff --git a/scripts/dev_evaluate_tool.py b/scripts/dev_evaluate_tool.py new file mode 100644 index 00000000..f29d327c --- /dev/null +++ b/scripts/dev_evaluate_tool.py @@ -0,0 +1,371 @@ +from __future__ import annotations + +import asyncio +import json +import os +import re +import sys +import time +from collections import defaultdict +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path + +from matplotlib import pyplot as plt +from tqdm import tqdm + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) + +os.environ["CLIMATECLAW_DEV"] = "1" # "1"/"true"/"yes" means dev mode +os.environ["CLIMATECLAW_LITE_LLM_ADDRESS"] = "http://localhost:4000" +os.environ["CLIMATECLAW_RAG_SERVER_URL"] = "http://localhost:8050" +os.environ["CLIMATECLAW_CODE_SERVER_URL"] = "http://localhost:8051" +os.environ["CLIMATECLAW_WEB_SEARCH_SERVER_URL"] = "http://localhost:8052" +os.environ["CLIMATECLAW_PLUGIN_CODE_SEARCH_SERVER_URL"] = "http://localhost:8053" +os.environ["CLIMATECLAW_MONGODB_HOST"] = "localhost" + +from climateclaw.core.logging_setup import configure_logging, silence_logger +from climateclaw.core.prompting import get_entire_prompt +from climateclaw.services.authentication.auth import Authenticator +from climateclaw.services.storage.helpers import create_dir_at_cache +from climateclaw.services.storage.mongodb_storage import ThreadStorage +from climateclaw.services.streaming.active_conversations import ( + end_and_save_conversation, + new_thread_id, +) +from climateclaw.services.streaming.stream_orchestrator import ( + prepare_for_stream, + run_stream, +) +from climateclaw.services.streaming.stream_variants import from_sv_to_json + +""" +Headless dev/benchmark runner mirroring /chatbot/streamresponse behaviour. Headless means that it runs without a web server, and without any user interaction. It is intended for development and benchmarking purposes. + +- Config below (no argparse). +- Disk-only persistence (never Mongo). +- RUNS/CONCURRENCY for benchmarks; set CONCURRENCY=1 for clean mode. +- Uses ONE global McpManager; orchestrator ties MCP session to thread_id. +""" + +# ────────────────────────────────────────────────────────────────────────────── +# CONFIG +# ────────────────────────────────────────────────────────────────────────────── + +MODEL = "gpt-4.1" # model to benchmark +USER_ID = "janedoe" +RUNS = 10 # number of runs to perform +CONCURRENCY = 10 # number of concurrent runs to perform + +PRINT_STREAM = False +PRINT_PER_RUN_SUMMARY = ( + False # print summary of each run (status, prompt, tool call, tool output) +) +SILENCE_LOGGING = True + +# run evaluation of tool calls (TP, FP, FN, accuracy, precision, recall) +EVAL_TOOL = True +PLOT_METRICS = True # plot metrics as bar charts at the end of the benchmark + +## evaluation suite: direct questions about a plugin + indirect ones (more widely phrased) +path_to_prompts = ( + Path(__file__).parent / "evaluation" / "benchmark-prompts-plugin-server.json" +) +with open(path_to_prompts, "r", encoding="utf-8") as f: + BENCHMARK = {k.lower(): v for k, v in json.load(f).items()} +# BENCHMARK = { +# "leadtimeselektor": [ +# "How does the 'leadtimeselektor' plugin work from a high-level perspective?", +# "How can I best extract lead times from decadal climate predictions?", +# "How are you?", +# ] +# } + +path_to_plugins = ( + Path(__file__).parents[1] + / "src/climateclaw/tools/plugin_code_search/available_plugins.md" +) +plugin_overview = path_to_plugins.read_text(encoding="utf-8") +plugins = re.findall(r"\*\*([^*]+)\*\*", plugin_overview) +ALL_PLUGINS = [p.lower() for p in plugins] + +# ────────────────────────────────────────────────────────────────────────────── +log = configure_logging("dev_script") + + +@dataclass +class RunResult: + idx: int + thread_id: str + # answer: str + tool_name: str + tool_args: str + tool_output: str + + +async def _run_once(idx: int, sem: asyncio.Semaphore, prompt: str) -> RunResult: + async with sem: + thread_id = await new_thread_id() + read_history = False + + Auth = Authenticator.local(username="dev_user") + Storage = await ThreadStorage.create() + create_dir_at_cache(thread_id) + + await prepare_for_stream( + thread_id=thread_id, + user_id=USER_ID, + Auth=Auth, + Storage=Storage, + read_history=read_history, + ) + + system_prompt = get_entire_prompt(USER_ID, thread_id, MODEL) + + t0 = time.perf_counter() + status = "Done" + # answer = "" + tool_name = "" + tool_args = "" + tool_output = "" + + try: + async for variant in run_stream( + model=MODEL, + thread_id=thread_id, + user_input=prompt, + system_prompt=system_prompt, + storage=Storage, + ): + # if getattr(variant, "variant", None) == "Assistant": + # txt_piece = getattr(variant, "text", "") or "" + # answer += txt_piece + + if getattr(variant, "variant", None) == "ToolCall": + tool_name = getattr(variant, "tool_name", "") + tool_args = getattr(variant, "arg", "") + + if getattr(variant, "variant", None) == "ToolOutput": + tool_output = getattr(variant, "output", "") + + if PRINT_STREAM and getattr(variant, "variant", None) != "Assistant": + print(json.dumps(from_sv_to_json(variant), ensure_ascii=False)) + + await end_and_save_conversation(thread_id, Storage) + + except Exception as e: + status = f"Error:{e}" + + duration = time.perf_counter() - t0 + + if PRINT_PER_RUN_SUMMARY: + print(f"Status: {status},\t Duration: {duration:.2f}s") + print(f"Prompt: {prompt}") + print(f"called tool: {tool_name}") + print(f" tool args: {tool_args}") + print(f" tool output: {tool_output}\n") + + return RunResult( + idx, + thread_id, + # answer, + tool_name, + tool_args, + tool_output, + ) + + +def _evaluate_tool_call_results( + tool_names: list[str], tool_args: list[str], tool_outputs: list[str], plugin: str +) -> dict[str, float]: + ## evaluate results w.r.t. whether the plugin tool was called at all by the model + true_pos_tool = sum( + [ + 1 + for tc, ta in zip(tool_names, tool_args) + if tc == "plugin_code_search" and "user_query" in ta + ] + ) + false_pos_tool = sum( + [1 for tc in tool_names if tc != "plugin_code_search" and tc] + ) # model called a plugin, but it was the wrong one + false_neg_tool = sum( + [1 for tc in tool_names if not tc] + ) # model did not call any plugin, but should have done so + + ## evaluate results w.r.t. whether the correct plugin was selected by the model + remaining_plugins = [p for p in ALL_PLUGINS if p != plugin] + correct_tool_calls = [ + True if tc == "plugin_code_search" else False for tc in tool_names + ] + true_pos_plugin = sum( + [1 for to, tc in zip(tool_outputs, correct_tool_calls) if plugin in to and tc] + ) # model called the correct plugin and got the correct output + false_pos_plugin = sum( + [ + 1 + for to, tc in zip(tool_outputs, correct_tool_calls) + if any([p in to for p in remaining_plugins]) and tc + ] + ) # model called a plugin, but it was the wrong one + false_neg_plugin = sum( + [ + 1 + for to, tc in zip(tool_outputs, correct_tool_calls) + if not any([p in to for p in ALL_PLUGINS]) and tc + ] + ) # model did not call any known plugin, but should have done so + + # compute metrics + accuracy_tool = true_pos_tool / RUNS + precision_tool = ( + true_pos_tool / (true_pos_tool + false_pos_tool) + if (true_pos_tool + false_pos_tool) > 0 + else 0 + ) + recall_tool = ( + true_pos_tool / (true_pos_tool + false_neg_tool) + if (true_pos_tool + false_neg_tool) > 0 + else 0 + ) + accuracy_plugin = true_pos_plugin / RUNS + precision_plugin = ( + true_pos_plugin / (true_pos_plugin + false_pos_plugin) + if (true_pos_plugin + false_pos_plugin) > 0 + else 0 + ) + recall_plugin = ( + true_pos_plugin / (true_pos_plugin + false_neg_plugin) + if (true_pos_plugin + false_neg_plugin) > 0 + else 0 + ) + + return { + "true_pos_tool": true_pos_tool / RUNS, + "false_pos_tool": false_pos_tool / RUNS, + "false_neg_tool": false_neg_tool / RUNS, + "true_pos_plugin": true_pos_plugin / RUNS, + "false_pos_plugin": false_pos_plugin / RUNS, + "false_neg_plugin": false_neg_plugin / RUNS, + "accuracy_tool": accuracy_tool, + "precision_tool": precision_tool, + "recall_tool": recall_tool, + "accuracy_plugin": accuracy_plugin, + "precision_plugin": precision_plugin, + "recall_plugin": recall_plugin, + } + + +async def run_benchmark_for_prompt( + prompt: str, plugin: str, eval_dict: dict[str, list[float]] +) -> dict[str, list[float]]: + + sem = asyncio.Semaphore(CONCURRENCY) + tasks = [asyncio.create_task(_run_once(i, sem, prompt)) for i in range(RUNS)] + results = await asyncio.gather(*tasks) + assert results, "No results returned from benchmark runs" + tool_names = [r.tool_name for r in results] + tool_args = [r.tool_args for r in results] + tool_outputs = [r.tool_output for r in results] + + if EVAL_TOOL: + metrics_dict = _evaluate_tool_call_results( + tool_names, tool_args, tool_outputs, plugin + ) + for k, v in metrics_dict.items(): + eval_dict[k].append(v) + + return eval_dict + + +def plot_metrics(avg_metrics: dict[str, float], save_dir: Path) -> None: + current_date = datetime.today().strftime("%Y-%m-%d") + # pie chart for tool/plugin TP, FP, FN + fig, axs = plt.subplots(1, 2, figsize=(12, 6)) + metrics = ["TP", "FP", "FN"] + tool_metrics = { + "True Positives": avg_metrics.get("true_pos_tool", 0), + "False Positives": avg_metrics.get("false_pos_tool", 0), + "False Negatives": avg_metrics.get("false_neg_tool", 0), + } + plugin_metrics = { + "True Positives": avg_metrics.get("true_pos_plugin", 0), + "False Positives": avg_metrics.get("false_pos_plugin", 0), + "False Negatives": avg_metrics.get("false_neg_plugin", 0), + } + charts = [tool_metrics, plugin_metrics] + titles = ["Tool Call", "Plugin Call"] + for ax, chart, title in zip(axs, charts, titles): + pie = ax.pie( + chart.values(), + autopct="%1.1f%%", + textprops=dict(fontsize=14, color="white", fontweight="bold"), + ) + ax.legend(pie.wedges, metrics, title="Metrics", loc="best", fontsize=14) + ax.set_title(title, fontsize=16) + plt.suptitle( + f"Evaluation Metrics for 'plugin_code_search' tool (Model: {MODEL}, Runs/Query: 30)" + ) + plt.tight_layout() + plt.subplots_adjust(wspace=0.4) + plt.savefig(save_dir / f"tool_evaluation_pie_charts_{current_date}.png", dpi=400) + + # bar chart for tool/plugin accuracy, precision, recall + fig, axs = plt.subplots(1, 2, figsize=(12, 6)) + for ax in axs: + ax.set_ylim(0, 1) + ax.set_ylabel("Metric Value", labelpad=12) + tool_metrics = { + "Accuracy": avg_metrics.get("accuracy_tool", 0), + "Precision": avg_metrics.get("precision_tool", 0), + "Recall": avg_metrics.get("recall_tool", 0), + } + plugin_metrics = { + "Accuracy": avg_metrics.get("accuracy_plugin", 0), + "Precision": avg_metrics.get("precision_plugin", 0), + "Recall": avg_metrics.get("recall_plugin", 0), + } + axs[0].bar(tool_metrics.keys(), tool_metrics.values()) + axs[0].set_title("Tool Metrics") + axs[1].bar(plugin_metrics.keys(), plugin_metrics.values()) + axs[1].set_title("Plugin Metrics") + plt.suptitle( + f"Evaluation Metrics for 'plugin_code_search' tool (Model: {MODEL}, Runs/Query: 30)" + ) + plt.tight_layout() + plt.subplots_adjust(wspace=0.3) + plt.savefig(save_dir / f"tool_evaluation_metrics_{current_date}.png", dpi=400) + + +async def main() -> None: + if SILENCE_LOGGING: + silence_logger() + + eval_dict = defaultdict(list) + prompt_count = sum(len(prompts) * RUNS for prompts in BENCHMARK.values()) + with tqdm(total=prompt_count, desc="Evaluation", unit="prompt") as progress: + for plugin, prompts in BENCHMARK.items(): + progress.set_postfix_str(f"plugin={plugin}") + for prompt in prompts: + # print(f"\n=== Running benchmark for prompt: {prompt} ===") + eval_dict = await run_benchmark_for_prompt(prompt, plugin, eval_dict) + progress.update(RUNS) + + # compute average metrics across all prompts + if EVAL_TOOL: + avg_metrics = {k: sum(v) / len(v) for k, v in eval_dict.items()} + print( + f"\n=== Average metrics [in %] across all prompts ({RUNS} runs/prompt) ===" + ) + for k, v in avg_metrics.items(): + print(f" {k:<17}: {v:.1%}") + + # plot metrics as bar charts (one for tool, one for plugin) + if PLOT_METRICS and EVAL_TOOL: + eval_dir = Path(__file__).parent / "evaluation" + eval_dir.mkdir(parents=True, exist_ok=True) + plot_metrics(avg_metrics, eval_dir) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/scripts/evaluation/benchmark-prompts-plugin-server.json b/scripts/evaluation/benchmark-prompts-plugin-server.json new file mode 100644 index 00000000..6539add7 --- /dev/null +++ b/scripts/evaluation/benchmark-prompts-plugin-server.json @@ -0,0 +1,52 @@ +{ + "leadtimeselektor": [ + "How does the 'leadtimeselektor' plugin work from a high-level perspective?", + "How can I best extract lead times from decadal climate predictions?", + "How can I select and aggregate specific lead times from decadal prediction ensembles?" + ], + "problems": [ + "How does the 'problems' plugin evaluate the predictive skill of a decadal climate experiment?", + "I want to perform skill score evaluation for climate data against observations. How can I do that?", + "How can I calculate climate prediction skill scores against observations or reanalysis data?" + ], + "psi": [ + "What precipitation indices can I compute using the 'psi' plugin?", + "How can I calculate precipitation indices from a precipitation time series?", + "How can I analyse precipitation characteristics with the 'psi' plugin?" + ], + "crops": [ + "What kind of crop impact analysis can I perform with the 'crops' plugin?", + "How can I assess the impact of extreme climate events on crop productivity?", + "How can I evaluate the effects of extreme and compound climate events on crops?" + ], + "hwmid": [ + "How does the 'hwmid' plugin work from a high-level perspective?", + "How can I evaluate heat waves using the Heat Wave Magnitude Index?", + "How can I calculate the Heat Wave Magnitude Index from climate data?" + ], + "climdexcalc2": [ + "What is the purpose of the 'climdexcalc2' plugin?", + "How can I calculate climate and extreme indices from daily temperature and precipitation data?", + "How can I derive climate indices from daily temperature and precipitation datasets?" + ], + "Regionmatch": [ + "How can I use 'Regionmatch' to find comparable districts based on climate and other regional characteristics?", + "How can I identify German districts that are climatically similar to a region of interest?", + "What criteria does the 'Regionmatch' plugin use to identify similar regions?" + ], + "Heat2UrbanImpact": [ + "How does the 'Heat2UrbanImpact' plugin work from a high-level perspective?", + "How can I identify heatwaves from regional climate model data for urban impact studies?", + "How can I prepare regional climate data for use in urban impact models?" + ], + "WayDown": [ + "How can I generate sub-hourly precipitation data from daily values with 'WayDown'?", + "How can I disaggregate daily precipitation data to the sub-hourly scale?", + "What is the purpose of the 'WayDown' plugin for urban hydrological impact studies?" + ], + "freva-plugin-template": [ + "I want to create a new Freva plugin. How can I do that?", + "What files and structure are required to create a plugin that is recognized by Freva?", + "How can I get started with developing a new Freva plugin using the template?" + ] +} diff --git a/src/climateclaw/core/prompting.py b/src/climateclaw/core/prompting.py index 4cc3553b..3ba099dd 100644 --- a/src/climateclaw/core/prompting.py +++ b/src/climateclaw/core/prompting.py @@ -24,8 +24,6 @@ Differences from Rust (documented for future parity) ---------------------------------------------------- -1) GPT-5: placeholder — we do NOT use GPT-5-specific prompt files yet; we log a warning - and fall back to the baseline prompt set. """ @@ -107,7 +105,7 @@ def _load_prompts( model: str, ) -> dict[Literal["starting", "summary", "examples_path"], str]: """ - Load raw prompt assets for the given model (with GPT-5 placeholder fallback). + Load raw prompt assets for the given model. Returns: { @@ -140,8 +138,6 @@ def _load_examples_as_messages(examples_path: str | Path) -> list[dict]: include_images=False, include_meta=True, # parity note: Rust typically drops meta; we keep for now ) # type: ignore[return-value] - # Note: OpenAIMessage is a class that inherits from TypedDict, so it can be used as a dict by json.dumps. - # This means that the above error message can be ignored. def get_entire_prompt(user_id: str, thread_id: str, model: str) -> list[dict[str, Any]]: @@ -154,13 +150,6 @@ def get_entire_prompt(user_id: str, thread_id: str, model: str) -> list[dict[str messages.append(_as_system_message(assets["starting"])) messages.extend(_load_examples_as_messages(assets["examples_path"])) messages.append(_as_system_message(assets["summary"])) - - # Optional: mark placeholder when model is GPT-5 (useful for debugging) - if model_is_gpt_5(model): - logger.info( - "GPT-5 placeholder active: baseline prompts used for model='%s'.", model - ) - return messages diff --git a/src/climateclaw/prompt_library/baseline/starting_prompt.md b/src/climateclaw/prompt_library/baseline/starting_prompt.md index 264f1d3a..2f8cfb76 100644 --- a/src/climateclaw/prompt_library/baseline/starting_prompt.md +++ b/src/climateclaw/prompt_library/baseline/starting_prompt.md @@ -99,6 +99,16 @@ Use `web_search` only to access online documentation related to: When answering with information from `web_search`, include inline citations with the URLs used. +### C.3. `plugin_code_search` (Plugin Code Lookup) + +Use this tool to extract relevant code from Freva plugins to **SUPPLEMENT** and **GUIDE** your answer or analysis in the context of the user's query. + +* Ground your explanation always in the retrieved plugin code. Follow a logical sequence, from high-level to detailed: purpose, input, output, and usage. +* For specific climate/weather questions, lay out the retrieved code *as reference (implementation)* to the user's query. +* End your response with the repo URLs of the relevant files (with `"levante"` as branch name). +* If no code was retrieved, the context is insufficient, or access was denied, summarize the returned message. +* Skip this tool for simple or generic operations (basic data loading, a single (zonal) mean/anomaly, a straightforward plot). + --- ## D. Data Access @@ -139,7 +149,7 @@ When answering with information from `web_search`, include inline citations with dset = xr.open_mfdataset(data_files) ``` -5. Always provide the databrowser host: +6. Always provide the databrowser host: ```python host = "nextgems.dkrz.de" @@ -239,7 +249,6 @@ These paths can be accessed directly. 9. When averaging geospatial or gridded quantities, consider whether area weighting is required. If grid cells represent different physical areas, use an area-weighted average rather than a simple arithmetic mean. - --- ## F. Plotting Standards diff --git a/src/climateclaw/prompt_library/baseline/summary_prompt.md b/src/climateclaw/prompt_library/baseline/summary_prompt.md index 8444d7f9..8cfe1292 100644 --- a/src/climateclaw/prompt_library/baseline/summary_prompt.md +++ b/src/climateclaw/prompt_library/baseline/summary_prompt.md @@ -2,7 +2,7 @@ ## Identity and Expertise -You are **ClimateClaw**, a helpful AI assistant at the German Climate Computing Center (**DKZ**). +You are **ClimateClaw**, a helpful AI assistant at the German Climate Computing Center (**DKRZ**). You specialize in climate and atmospheric data analysis, particularly reanalysis and model data. Your capabilities include: @@ -47,6 +47,12 @@ For questions about the following topics, use the `web_search` tool to consult o When using `web_search`, always provide inline citations containing the URLs that were used. +## Plugin Code Search + +Use the `plugin_code_search` tool for repository-grounded code knowledge that **SUPPLEMENTS** and **GUIDES** the standard routine. +Call it for either direct questions regarding plugins or for specific climate/weather questions that could be answered by a plugin. +When requested, provide a lightweight code snippet based on the returned context. + ## Required Workflow For tasks requiring analysis, data processing, or code execution: diff --git a/src/climateclaw/prompt_library/local_model/starting_prompt.md b/src/climateclaw/prompt_library/local_model/starting_prompt.md index 4f1a74cf..8fe5b7a7 100644 --- a/src/climateclaw/prompt_library/local_model/starting_prompt.md +++ b/src/climateclaw/prompt_library/local_model/starting_prompt.md @@ -112,6 +112,15 @@ Use `web_search` only to access online documentation related to: When answering with information from `web_search`, include inline citations with URLs. +### C.3. `plugin_code_search` (Plugin Code Lookup) + +Use the retrieved Freva plugin code to **SUPPLEMENT** and **GUIDE** your answer or analysis whenever an established plugin covers the user's task. Ground your explanation in the retrieved plugin code. Do not speculate about the plugin's internal logic or implementation details. + +* Explain the plugin's logic and how to use it; for specific climate/weather questions, explain how it *relates as reference* to the user's query. +* End your response with the repo URLs of the relevant files (with `"levante"` as branch name). +* If no code was retrieved, the context is insufficient, or access was denied, say so explicitly (summarizing the returned message) and ask for more details or suggest checking GitLab access rights. +* Skip this tool for simple or generic operations (basic data loading, a single (zonal) mean/anomaly, a straightforward plot) with no specialized methodology involved. + --- ## D. Data Access diff --git a/src/climateclaw/prompt_library/local_model/summary_prompt.md b/src/climateclaw/prompt_library/local_model/summary_prompt.md index 8444d7f9..86fb871f 100644 --- a/src/climateclaw/prompt_library/local_model/summary_prompt.md +++ b/src/climateclaw/prompt_library/local_model/summary_prompt.md @@ -47,6 +47,14 @@ For questions about the following topics, use the `web_search` tool to consult o When using `web_search`, always provide inline citations containing the URLs that were used. +## Plugin-Code + +Use the `plugin_code_search` tool for repository-grounded code knowledge that **SUPPLEMENTS** and **GUIDES** the standard routine (*load data → compute → plot*). Call it: +1. when the user directly asks about a specific Freva plugin's logic, usage, or configuration; or +2. proactively, deciding on your own, for more complex regional, decadal, or extreme-event analysis questions where a fitting plugin repository likely contains the domain-specific methodology — call it first to ground the scripting rather than improvising the method. + +Skip it for simple, generic operations already covered by the standard workflow. When requested, use `code_interpreter` to write working code based on the returned context. + ## Required Workflow For tasks requiring analysis, data processing, or code execution: diff --git a/src/climateclaw/services/mcp/mcp_manager.py b/src/climateclaw/services/mcp/mcp_manager.py index e86d8279..18dd59e0 100644 --- a/src/climateclaw/services/mcp/mcp_manager.py +++ b/src/climateclaw/services/mcp/mcp_manager.py @@ -4,6 +4,7 @@ import os from typing import Any, Literal +from climateclaw.core.available_chatbots import available_chatbots, default_chatbot from climateclaw.core.logging_setup import configure_logging from climateclaw.core.settings import get_settings from climateclaw.services.authentication.auth import Authenticator @@ -234,6 +235,8 @@ def get_mcp_headers( ) -> dict[str, dict[str, str | None]]: mongodb_uri = get_mongodb_uri() + model = "gpt-5.6-luna" + headers: dict[str, dict[str, str | None]] = { "rag-server": { "mongodb-uri": mongodb_uri, @@ -241,5 +244,9 @@ def get_mcp_headers( "code-server": { "working-dir": str(cache), }, + "plugin-code-search-server": { + "username": auth.username or "unknown_user", + "model": model if model in available_chatbots() else default_chatbot(), + }, } return headers diff --git a/src/climateclaw/services/storage/summarize_topic.py b/src/climateclaw/services/storage/summarize_topic.py index f88e9dc1..d57c6271 100644 --- a/src/climateclaw/services/storage/summarize_topic.py +++ b/src/climateclaw/services/storage/summarize_topic.py @@ -3,7 +3,7 @@ import re from typing import List -from climateclaw.core.available_chatbots import default_chatbot_local +from climateclaw.core.available_chatbots import default_chatbot from climateclaw.core.logging_setup import configure_logging from climateclaw.services.streaming.litellm_client import acomplete, first_text from climateclaw.services.streaming.stream_variants import StreamVariant, SVUser @@ -101,7 +101,8 @@ async def summarize_topic(content: List[StreamVariant]) -> str: try: resp = await acomplete( messages=[{"role": "user", "content": prompt}], - model=default_chatbot_local(), + # model=default_chatbot_local(), + model=default_chatbot(), max_tokens=30, temperature=0.1, ) diff --git a/src/climateclaw/services/streaming/litellm_client.py b/src/climateclaw/services/streaming/litellm_client.py index 97ca1cca..3b8e9c9a 100644 --- a/src/climateclaw/services/streaming/litellm_client.py +++ b/src/climateclaw/services/streaming/litellm_client.py @@ -69,7 +69,7 @@ async def acomplete( stream: bool = False, temperature: float | None = None, max_tokens: int | None = None, - extra: dict[str, Any] | None = None, + extra_params: dict[str, Any] | None = None, **request_params: Any, ) -> dict[str, Any] | AsyncIterator[dict[str, Any]]: """ @@ -87,8 +87,8 @@ async def acomplete( payload["temperature"] = temperature if max_tokens is not None: payload["max_tokens"] = max_tokens - if extra: - payload.update(extra) + if extra_params: + payload.update(extra_params) payload.update(_passthrough_params(None)) if request_params: payload.update(_passthrough_params(request_params)) diff --git a/src/climateclaw/services/streaming/stream_orchestrator.py b/src/climateclaw/services/streaming/stream_orchestrator.py index 74572375..dcf60a66 100644 --- a/src/climateclaw/services/streaming/stream_orchestrator.py +++ b/src/climateclaw/services/streaming/stream_orchestrator.py @@ -449,7 +449,10 @@ async def run_with_heartbeat(): ( tool_out_v, tool_msgs, - ) = r.var_block, r.tool_messages + ) = ( + r.var_block, + r.tool_messages, + ) break else: yield r # Streaming the result to endpoint diff --git a/src/climateclaw/services/streaming/tool_calls.py b/src/climateclaw/services/streaming/tool_calls.py index b478507c..94e27124 100644 --- a/src/climateclaw/services/streaming/tool_calls.py +++ b/src/climateclaw/services/streaming/tool_calls.py @@ -95,12 +95,12 @@ def accumulate_tool_calls(delta: Dict[str, Any], agg: Dict[str, Any]) -> None: ) if item.get("id"): entry["id"] = item["id"] - f = item.get("function") or {} - if f.get("name"): - entry["function"]["name"] = f["name"] - if f.get("arguments"): + fn = item.get("function") or {} + if fn.get("name"): + entry["function"]["name"] = fn["name"] + if fn.get("arguments"): entry["function"]["arguments"] = ( - entry["function"].get("arguments", "") + f["arguments"] + entry["function"].get("arguments", "") + fn["arguments"] ) @@ -425,4 +425,5 @@ def parse_generic_tool_result(result: dict, tool_name: str, id: str, logger=None out = "Unknown response." web_sv = SVToolOutput(content=out, tool_name=tool_name, id=id) # type: ignore[arg-type] web_msg = help_convert_sv_ccrm([web_sv]) + yield web_sv yield FinalSummary(var_block=[web_sv], tool_messages=web_msg, is_error=False) diff --git a/src/climateclaw/tools/plugin_code_search/available_plugins.md b/src/climateclaw/tools/plugin_code_search/available_plugins.md new file mode 100644 index 00000000..cb02f3ba --- /dev/null +++ b/src/climateclaw/tools/plugin_code_search/available_plugins.md @@ -0,0 +1,44 @@ +# Freva Plugins + +Below is a list of available plugins for decadal prediction skill assessment and related tasks. +Each of the following sections belongs to a specific Freva instance/project, where each entry in it includes + +* the plugin name; and +* a brief description of its functionality. + +## "coming decade" + +* **leadtimeselektor**: extracts and aggregates lead times from decadal prediction ensembles +* **problems**: performs skill score calculation of a decadal climate experiment (and a reference experiment, if applicable) against reanalysis or observation +* **cvprepare**: prepares cross-validation datasets for decadal prediction skill assessment +* **recalibration**: re-calibrates decadal climate data to observations for correction of model drift and bias +* **terciles**: computes tercile-based statistics for prediction skill assessment + +--- + +## "climxtreme" + +* **climdexcalc2**: calculates climatological & extreme indices from daily temperature and precipitation data +* **preproc**: pre-processes spatio-temporal data, converting to CMORized formats +* **crops**: assesses the impact of extreme and compound climate events on crop productivity +* **hwmid**: heat wave evaluation using Heat Wave Magnitude Index daily +* **psi**: calculates different precipitation indices for a given precipitation time series +* **precip_return_period_maps**: computes maps of precipitation sums and return periods +* **pca**: performs principal component analysis / Empirical Orthogonal Functions for a given spatio-temporal 3D dataset + +--- + +## "regiklim" + +* **Climpact**: processes climate model output data for usage as input of impact model studies. +* **Heat2UrbanImpact**: identifies heatwaves from regional climate model data for usage in urban impact models +* **IDF-CCF**: calculates intensity-duration-frequency, climate change factors +* **Regionmatch**: identifies German districts similar to a region of interest based on climate, sociodemography, orography and land use, supporting regional adaptation planning and knowledge transfer +* **VisualizeClimDexCalc**: visualization of climate indices +* **WayDown**: Temporal disaggregation of daily precipitation to sub-hourly scale with a markov chain model for urban hydrological impact studies + +--- + +## "freva" + +* **freva-plugin-template**: defines a template for creating new Freva plugins, including the necessary structure and files for a plugin to be recognized by the Freva system. diff --git a/src/climateclaw/tools/plugin_code_search/server.py b/src/climateclaw/tools/plugin_code_search/server.py new file mode 100644 index 00000000..11ee8f4f --- /dev/null +++ b/src/climateclaw/tools/plugin_code_search/server.py @@ -0,0 +1,568 @@ +import ast +import json +import os +import re +import textwrap +from contextvars import ContextVar +from pathlib import Path +from typing import Optional, Tuple +from urllib.parse import quote as urlquote + +import httpx +from fastmcp import FastMCP + +from climateclaw.core.logging_setup import configure_logging +from climateclaw.services.streaming.litellm_client import acomplete, first_text +from climateclaw.tools.header_gate import RequestIdMiddleware, make_header_gate + +logger = configure_logging(__name__) + + +# ── Tool Config ──────────────────────────────────────────────────────────────── +GITLAB_ACCESS_TOKEN: Optional[str] = os.getenv("CLIMATECLAW_GITLAB_ACCESS_TOKEN") +GITLAB_BASE_URL = "https://gitlab.dkrz.de/api/v4" +FREVA_PROJECT_NAMES = { + "coming decade": "kd1418", + "climxtreme": "bm1159", + "regiklim": "ch1187", + "freva": "freva", +} +ALLOWED_FILE_EXTENSIONS = ( + ".py", + ".sh", + ".R", + ".md", + ".rst", +) # only fetch these file types +MAX_FILE_SIZE_BYTES = 50_000 # skip files larger than this +MAX_TOTAL_CODE_CHARS = 100_000 # truncate total fetched code after this limit +MAX_FILES = 6 # max files the LLM may select in total +ENTRY_PATTERN = re.compile(r"(wrapper|api)[^/]*\.py$", re.IGNORECASE) +EXCLUDE_PATTERN = re.compile(r"(^|/)(tests?/|test_|__init__\.py$)", re.IGNORECASE) + + +# ── Server Config ──────────────────────────────────────────────────────────────── +HOST = os.getenv("CLIMATECLAW_MCP_HOST", "0.0.0.0") +PORT = int(os.getenv("CLIMATECLAW_MCP_PORT", "8053")) +PATH = os.getenv("CLIMATECLAW_MCP_PATH", "/mcp") # standard path +USERNAME = "username" +MODEL = "model" + + +# ─── App ──────────────────────────────────────────────────────────────────── +mcp = FastMCP("plugin-code-search-server") +mcp.add_middleware(RequestIdMiddleware()) +logger.info("Starting Freva-Plugin Code Search MCP server on %s:%s%s", HOST, PORT, PATH) + +# Per-request header context +username_ctx: ContextVar[str | None] = ContextVar("username_ctx", default=None) +model_ctx: ContextVar[str | None] = ContextVar("model_ctx", default=None) + +# Start the MCP server using Streamable HTTP transport +app = make_header_gate( + mcp.http_app(), + ctx_list=[username_ctx, model_ctx], + header_name_list=[USERNAME, MODEL], + logger=logger, + mcp_path=PATH, +) + + +def _get_user(): + user = username_ctx.get() + if not user: + logger.warning(f"Missing required header '{USERNAME}'! ") + return "unknown_user" + else: + return user + + +def _get_model(): + model = model_ctx.get() + if not model: + logger.warning(f"Missing required header '{MODEL}'! ") + return MODEL + else: + return model + + +async def call_llm(prompt: str) -> str: + model = _get_model() + resp = await acomplete( + model=model, + messages=[{"role": "user", "content": prompt}], + ) + raw_text = first_text(resp).strip() + return raw_text + + +# ── GitLab helpers ──────────────────────────────────────────────── +_gitlab_http = httpx.Client( + base_url=GITLAB_BASE_URL, + headers={"PRIVATE-TOKEN": GITLAB_ACCESS_TOKEN or ""}, + timeout=30.0, +) + + +def _fetch_file_raw(project_id: int, file_path: str, branch: str) -> str: + encoded_path = urlquote(file_path, safe="") + resp = _gitlab_http.get( + f"/projects/{project_id}/repository/files/{encoded_path}/raw", + params={"ref": branch}, + ) + resp.raise_for_status() + return resp.text + + +def _has_read_access(project_id: int, username: str) -> bool: + """ + Check if the given user has at least read access rights to the GitLab project, + based on the project's visibility, the user's ID as well as membership status. + """ + # Get project visibility: "public", "internal", or "private" + resp = _gitlab_http.get(f"/projects/{project_id}") + resp.raise_for_status() + visibility = resp.json().get("visibility") + # Public: everyone can read + if visibility == "public": + return True + + # Internal: all authenticated (non-external) users can read + resp = _gitlab_http.get("/users", params={"username": username}) + resp.raise_for_status() + users = resp.json() + user_id = users[0].get("id") if users else None + if user_id is None: + return False + + if visibility == "internal": + return True + + # Private: must have explicit membership + resp = _gitlab_http.get(f"/projects/{project_id}/members/all/{user_id}") + if resp.status_code == 404: + return False + resp.raise_for_status() + return resp.json().get("access_level") >= 10 # Guest and above can read + + +def get_project_id(plugin: str, project: str) -> int | None: + """Fetch the GitLab project ID for the given plugin name.""" + encoded = urlquote(f"{project}/plugins4freva/{plugin}", safe="") + resp = _gitlab_http.get(f"/projects/{encoded}") + if resp.status_code == 404: + return None + resp.raise_for_status() + return resp.json().get("id") + + +def fetch_repo_tree_and_branch(project_id: int) -> tuple[list[str], str]: + """ + Return the recursive file tree for a plugin repository as a list of + file paths (strings) that match relevant extensions. + """ + items: list[dict] = [] + page = 1 + while True: + resp = _gitlab_http.get( + f"/projects/{project_id}/repository/tree", + params={"recursive": "true", "per_page": 100, "page": page}, + ) + resp.raise_for_status() + batch = resp.json() + if not batch: + break + items.extend(batch) + page += 1 + + # filter to file paths with relevant extensions + paths = [ + entry["path"] + for entry in items + if entry["type"] == "blob" and entry["path"].endswith(ALLOWED_FILE_EXTENSIONS) + ] + + # Get the default branch of the repository (usually "levante" or "master") + resp = _gitlab_http.get(f"/projects/{project_id}") + resp.raise_for_status() + default_branch = resp.json().get("default_branch") + return paths, default_branch + + +def fetch_files( + project_id: int, branch: str, selected_files: list[str], max_chars: int +) -> dict[str, str]: + """ + Fetch the raw content of selected files until the total character count + reaches `max_chars`. Returns a mapping of file path -> content. + """ + collected: dict[str, str] = {} + budget = max_chars + for file in selected_files: + if budget <= 0: + break + try: + content = _fetch_file_raw(project_id, file, branch) + except Exception as e: + logger.debug( + "Skipping file %s due to error in fetching content: %s", file, e + ) + continue + if len(content) > MAX_FILE_SIZE_BYTES: + content = content[:MAX_FILE_SIZE_BYTES] + "\n... (file truncated)" + if len(content) > budget: + content = ( + content[:budget] + f"\n... (truncated: reached {max_chars} char limit)" + ) + collected[file] = content + budget -= len(content) + return collected + + +# ── Other helpers ──────────────────────────────────────────────── + + +def extract_method_from_source(source_code: str, method_name: str = "run_tool") -> str: + """Parses a Python source string and returns the code of a specific method.""" + try: + # Clean tabs and fix uneven leading indentation + cleaned_code = source_code.expandtabs(4) + cleaned_code = textwrap.dedent(cleaned_code) + # Parse the cleaned source code + tree = ast.parse(cleaned_code) + # Look for the method definition anywhere in the tree + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == method_name: + method_code = ast.get_source_segment(source_code, node) + return method_code or "No 'run_tool' method found in the entry file." + + return f"Method '{method_name}' not found in the source code." + + except SyntaxError as e: + return f"Syntax error in the provided source code: {e}" + + +def format_files(files: dict[str, str]) -> str: + """Concatenate fetched files into a single string, separated by file headers.""" + return "\n".join( + f"### FILE: {path} ###\n```\n{content}\n```\n" + for path, content in files.items() + ) + + +# ────────────────────────────────────────────────────────────────────────── + + +def find_entry_file(file_paths: list[str]) -> str: + """ + Deterministically find the plugin entry file ('*wrapper*.py' or '*api*.py'), + excluding tests. Shallower paths come first, as they are most likely the + actual plugin entry point. + """ + matches = [ + p + for p in file_paths + if ENTRY_PATTERN.search(p) and not EXCLUDE_PATTERN.search(p) + ] + return sorted(matches, key=lambda p: (p.count("/"), p))[0] + + +async def _llm_select_files( + prompt: str, file_paths: list[str], num_files: int, fallback: list[str] +) -> list[str]: + """ + Ask the LLM for a JSON array of file paths and keep only those that exist + in `file_paths`. Returns `fallback` if the LLM response cannot be parsed. + """ + try: + raw_text = await call_llm(prompt) + # Strip markdown code fences if present + if raw_text.startswith("```"): + raw_text = "\n".join(raw_text.split("\n")[1:]) + if raw_text.endswith("```"): + raw_text = "\n".join(raw_text.split("\n")[:-1]) + select_files: list[str] = json.loads(raw_text.strip()) + except Exception as e: + logger.warning("LLM file selection failed (%s); falling back to heuristic.", e) + select_files = fallback + + # Return only paths that exist in repo tree and conform to exclusion pattern + valid_paths = [ + p + for p in select_files + if p in set(file_paths) and not EXCLUDE_PATTERN.search(p) + ] + return valid_paths[:num_files] # limit to max number of files + + +async def select_init_files( + plugin: str, + user_query: str, + fetched_code: str, + file_paths: list[str], + num_files: int, +) -> list[str]: + """ + Fallback if no entry file was found: let the LLM pick the most relevant files + from the repo tree for the user's query. + """ + file_tree = "\n".join(file_paths) + selection_prompt = ( + f"Task: You are analyzing the files from the '{plugin}' Freva plugin repository, given the user context and the fetched code. " + "Prioritize and select the files based on the following rules:\n\n" + "- Select the files that are imported or called as scripts in the fetched code, as well as those that seem *most* relevant to answer the user's query.\n" + "- For high level usage/configuration questions, prioritize config files, README and docs.\n" + "- For questions about implementation logic, prioritize source code modules.\n" + "- Exclude tests, examples, generated files, and any '__init__.py'.\n" + f"- Return ONLY a valid JSON array of at most {num_files} file path strings from the provided list. Output nothing but the JSON array.\n\n" + f"=== User Query ===\n{user_query}\n=== END ===\n\n" + f"=== Repository file list ===\n{file_tree}\n=== END ===\n\n" + f"=== Fetched Code ===\n{fetched_code}\n=== END ===\n\n" + ) + # Fallback: pick all files from the tree search + return await _llm_select_files( + selection_prompt, file_paths, num_files, fallback=file_paths + ) + + +async def select_dependency_files( + plugin: str, + user_query: str, + fetched_code: str, + file_paths: list[str], + num_files: int, +) -> list[str]: + """ + Ask the LLM which of the remaining repo files are direct dependencies of the + fetched code. + """ + file_tree = "\n".join(file_paths) + selection_prompt = ( + f"Task: You are tracing the code base of the '{plugin}' Freva plugin, given the user query and the already fetched source code. " + "Prioritize and select which of the remaining repository files should be read next based on the following rules:\n\n" + "- Scan the fetched code for imports or script calls and map each to the matching file using Python module conventions " + "(e.g., 'from foo.bar import baz' -> 'foo/bar.py').\n" + "- Ignore imports from third-party libraries (e.g. numpy, xarray, evaluation_system etc.) and focus on the plugin's own source code.\n" + "- Exclude tests, examples, generated files, and any '__init__.py'.\n" + f"- Return ONLY a valid JSON array of at most {num_files} file path strings from the remaining list. If none match, return []. Output nothing but the JSON array.\n\n" + f"=== User Query ===\n{user_query}\n=== END ===\n\n" + f"=== Fetched Code ===\n{fetched_code}\n=== END ===\n\n" + f"=== Remaining Repository Files ===\n{file_tree}\n=== END ===\n\n" + ) + # Fallback: pick nothing to avoid fetching irrelevant code + return await _llm_select_files(selection_prompt, file_paths, num_files, fallback=[]) + + +async def collect_plugin_context( + plugin: str, project: str, project_id: int, user_query: str +) -> str: + """ + Entry-point-driven context retrieval of the plugin code base: + 1. Deterministically find the plugin entry files ('*wrapper*.py' / '*api*.py') + in the repo tree. + 2. Let the LLM select the direct dependencies of the entry files plus other + modules relevant for the user's query, and fetch those. + 3. Let the LLM resolve the direct dependencies of the newly fetched files. + Returns a string containing the concatenated relevant source code files, + separated by file and with a header. + """ + + def _log_stage(stage: str, files: list[str]): + logger.info( + "%s retrieval step selected %d/%d files for plugin '%s': %s", + stage, + len(files), + len(file_paths), + plugin, + files, + ) + + header = ( + f"Relevant retrieved code of the '{plugin}' plugin " + f"(https://gitlab.dkrz.de/{project}/plugins4freva/{plugin}):\n\n" + ) + + # ── Fetch the repository tree with all files ──────────────────── + file_paths, branch = fetch_repo_tree_and_branch(project_id) + if not file_paths: + return f"repository is empty for plugin '{plugin}' in branch '{branch}'" + + # ── Stage 1: find & fetch plugin entry ───────────────────── + entry_file = find_entry_file(file_paths) + _log_stage("Deterministic entry file", [entry_file]) + + exec_code = extract_method_from_source( + _fetch_file_raw(project_id, entry_file, branch), method_name="run_tool" + ) + entry_code = format_files({entry_file: exec_code}) + + # ── Stage 2: find & fetch useful files, based on entry file + user context ─────── + remaining = [p for p in file_paths if p != entry_file] + init_files = await select_init_files( + plugin, user_query, entry_code, remaining, MAX_FILES // 2 + ) + init_files = [entry_file] + init_files + init_code = fetch_files(project_id, branch, init_files, MAX_TOTAL_CODE_CHARS) + _log_stage("Combined LLM-based initial", init_files) + + # ── Stage 3: resolve dependencies ─────────────────────────────────── + budget = MAX_TOTAL_CODE_CHARS - sum(len(c) for c in init_code.values()) + remaining = [p for p in file_paths if p not in init_code] + max_deps = MAX_FILES - len(init_files) + dep_files = await select_dependency_files( + plugin, + user_query, + format_files(init_code), + remaining, + max_deps, + ) + _log_stage("LLM-based dependency", dep_files) + deps_code = fetch_files(project_id, branch, dep_files, budget) + + # Format the final output with header and fetched code, including dependencies + code_content = format_files(init_code) + if deps_code: + code_content += "\n\n### ── Dependency files ── ###\n\n" + format_files( + deps_code + ) + return header + code_content + + +def validate_plugin_call( + plugin: str, project: str, project_id: int | None +) -> Tuple[bool, str]: + """ + Validate GitLab repo availability for chosen plugin / project names, + as well as reading access for the current user. + Returns a tuple (bool, str): + - False and an error message string if validation fails; or + - True and a success message if valid. + """ + # validate GitLab access of user + if project_id is None: + error_msg = f"Plugin '{plugin}' not found in GitLab project '{project}'." + return False, error_msg + + # TODO: Username hardcoded for now; replace with _get_user() before merge/deployment + user_name = "k202218" + # user_name = _get_user() + try: + user_access = _has_read_access(project_id, user_name) + if not user_access: + logger.warning( + "User '%s' does NOT have read access to plugin '%s' in project '%s'.", + user_name, + plugin, + project, + ) + error_msg = ( + f"User access for {user_name} to plugin '{plugin}' denied! " + f"Get access by being added to GitLab project '{project}'." + ) + return False, error_msg + logger.info( + "Authorization layer passed: User has read access to plugin '%s' in project '%s'.", + plugin, + project, + ) + except httpx.HTTPError as e: + logger.error("Error checking GitLab repo membership: %s", e) + error_msg = "Plugin code search is currently unavailable (GitLab access error)." + return False, error_msg + + return True, "Validation successful" + + +async def detect_plugin_project(user_query: str) -> tuple[str, str]: + """ + Attempt to automatically detect the plugin and project names from the user's query + by making a call to the LLM. + + Returns a tuple of (plugin_name, project_name): + - plugin_name (str): Name of the repo or plugin (e.g. "leadtimeselektor") + - project_name (str): Name of the Freva instance (e.g. "kd1418" for "coming decade") + """ + file_path = Path(__file__).parent / "available_plugins.md" + plugin_descriptions = file_path.read_text(encoding="utf-8", newline="\n") + # let LLM select plugin and project names from the available plugin summaries + selection_prompt = ( + "Task: Select the semantically best matching Freva plugin & project name for the user query from the provided plugin summaries.\n" + "Rules:\n" + "- Select only plugin & project names that appear in the summaries.\n" + "- Return exactly one plugin and its corresponding project.\n" + "- Output must be exactly this format with no extra text: ,\n\n" + f"User query:\n{user_query}\n\n" + f"Available plugin summaries:\n{plugin_descriptions}" + ) + raw_text = await call_llm(selection_prompt) + # format to "," output + plugin_name, project_name = raw_text.split(",", maxsplit=1) + plugin_name = plugin_name.strip().lower() + project_name = FREVA_PROJECT_NAMES.get(project_name.strip().lower(), "") + logger.info( + "Auto-detected plugin '%s' and project '%s' for examining code base.", + plugin_name, + project_name, + ) + return plugin_name, project_name + + +@mcp.tool() +async def plugin_code_search(user_query: str) -> str: + """ + Fetch relevant source code and documentation files of a Freva data analysis plugin + as a repository-grounded code knowledge base. + + Use this when the user + - explicitly asks how a plugin's internal logic works, how to run or configure it, + or wants plugin code translated or adapted into Python examples; + - asks a specific climate or weather question involving regional, decadal, or + extreme-event analysis (lead time selection, hindcast skill scoring, bias adjustment + and drift correction, downscaling, precipitation indices, crop impact, heat waves/HWMID + climate/extreme indices, region matching, urban heatwave extraction, precipitation + disaggregation, and plugin creation). + + Workflow Guidelines: + - Ground explanations only in returned code context. + - Handle retrieved plugin code in two separate steps: + 1. **First step (always) – only high level:** a factful explanation of how the plugin works and how to use it (including plugin and project names). + 2. **Second step (only if requested) – implementation:** turn the plugin's *core logic* into a lightweight Python snippet with a concise plan. + - reference modules, classes, and functions only when asked for more detail. + - call it again for follow-up questions *not* sufficiently covered by prior results. + - only implement the plugin's *core logic* in a lightweight Python snippet - for that: + - replace `cdo` with `xarray` operations + - prioritize workflow correctness over mirroring non-critical details (e.g. multiple tries, fallbacks, logging) + + Args: + ----- + user_query (str): What the user wants to know about or do with a dedicated Freva + plugin regarding climate or weather data analysis. Always the only argument. + + Returns: + -------- + str: Relevant code context (with a header containing the plugin's repo URL), + including directly imported dependency files; or an error message if the plugin + is not found, user access is denied, or code retrieval fails. + + Examples: + --------- + - "How does the 'leadtimeselektor' plugin work from a high-level perspective?" + - "How can I calculate climate prediction skill scores against observations or reanalysis data?" + - "How can I assess the impact of extreme climate events on crop productivity?" + """ + plugin, project = await detect_plugin_project(user_query) + # return f"plugin_name: {plugin}, project_name: {project}" # only for benchmarking the plugin detection stage + + # Validate the plugin call + project_id = get_project_id(plugin, project) + result, message = validate_plugin_call(plugin, project, project_id) + if not result: + return message + + # Fetch the plugin code and return it with a header + logger.info("Fetching source code for plugin '%s'", plugin) + try: + return await collect_plugin_context(plugin, project, project_id, user_query) # type: ignore + except Exception as e: + logger.warning("Failed to fetch plugin code for '%s': %s", plugin, e) + return f"Failed to retrieve source code for plugin '{plugin}': {e}"