seraph-mcp 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (43) hide show
  1. seraph_mcp-0.1.0/.gitignore +10 -0
  2. seraph_mcp-0.1.0/.python-version +1 -0
  3. seraph_mcp-0.1.0/AGENT_SETUP.md +45 -0
  4. seraph_mcp-0.1.0/LICENSE +21 -0
  5. seraph_mcp-0.1.0/PKG-INFO +115 -0
  6. seraph_mcp-0.1.0/README.md +56 -0
  7. seraph_mcp-0.1.0/pyproject.toml +86 -0
  8. seraph_mcp-0.1.0/requirements.txt +1 -0
  9. seraph_mcp-0.1.0/src/seraph/__init__.py +5 -0
  10. seraph_mcp-0.1.0/src/seraph/backends/__init__.py +0 -0
  11. seraph_mcp-0.1.0/src/seraph/backends/common.py +102 -0
  12. seraph_mcp-0.1.0/src/seraph/backends/gemini.py +186 -0
  13. seraph_mcp-0.1.0/src/seraph/cli.py +162 -0
  14. seraph_mcp-0.1.0/src/seraph/config.py +281 -0
  15. seraph_mcp-0.1.0/src/seraph/graph/__init__.py +0 -0
  16. seraph_mcp-0.1.0/src/seraph/graph/builder.py +374 -0
  17. seraph_mcp-0.1.0/src/seraph/graph/traversal.py +149 -0
  18. seraph_mcp-0.1.0/src/seraph/index.py +643 -0
  19. seraph_mcp-0.1.0/src/seraph/lineage.py +178 -0
  20. seraph_mcp-0.1.0/src/seraph/memory.py +135 -0
  21. seraph_mcp-0.1.0/src/seraph/protocols.py +79 -0
  22. seraph_mcp-0.1.0/src/seraph/query/__init__.py +0 -0
  23. seraph_mcp-0.1.0/src/seraph/query/analyzer.py +65 -0
  24. seraph_mcp-0.1.0/src/seraph/query/classifier.py +86 -0
  25. seraph_mcp-0.1.0/src/seraph/query/decomposition.py +35 -0
  26. seraph_mcp-0.1.0/src/seraph/query/llm.py +84 -0
  27. seraph_mcp-0.1.0/src/seraph/query/parser.py +171 -0
  28. seraph_mcp-0.1.0/src/seraph/query/weights.py +206 -0
  29. seraph_mcp-0.1.0/src/seraph/retrieval/__init__.py +0 -0
  30. seraph_mcp-0.1.0/src/seraph/retrieval/base.py +22 -0
  31. seraph_mcp-0.1.0/src/seraph/retrieval/dedup.py +150 -0
  32. seraph_mcp-0.1.0/src/seraph/retrieval/fusion.py +91 -0
  33. seraph_mcp-0.1.0/src/seraph/retrieval/lexical.py +70 -0
  34. seraph_mcp-0.1.0/src/seraph/retrieval/pipeline.py +345 -0
  35. seraph_mcp-0.1.0/src/seraph/retrieval/reranker.py +190 -0
  36. seraph_mcp-0.1.0/src/seraph/retrieval/semantic.py +359 -0
  37. seraph_mcp-0.1.0/src/seraph/retrieval/structural.py +311 -0
  38. seraph_mcp-0.1.0/src/seraph/retrieval/tokenize.py +41 -0
  39. seraph_mcp-0.1.0/src/seraph/retrieval/vectors.py +50 -0
  40. seraph_mcp-0.1.0/src/seraph/server.py +118 -0
  41. seraph_mcp-0.1.0/src/seraph/service.py +284 -0
  42. seraph_mcp-0.1.0/src/seraph/smoke.py +110 -0
  43. seraph_mcp-0.1.0/src/seraph/types.py +183 -0
@@ -0,0 +1,10 @@
1
+ __pycache__/
2
+ *.pyc
3
+ .venv/
4
+ .pytest_cache/
5
+ .ruff_cache/
6
+ .seraph/
7
+ .seraph_cache/
8
+ *.egg-info/
9
+ dist/
10
+ build/
@@ -0,0 +1 @@
1
+ 3.12
@@ -0,0 +1,45 @@
1
+ # Set up Seraph in this coding agent
2
+
3
+ Use these instructions when a user asks to install or use Seraph. Seraph is a local MCP server for searching local Git repositories. MCP works across many coding agents, but each agent has its own server registration mechanism.
4
+
5
+ ## Install
6
+
7
+ Requirements: Git, Python 3.12, and `uv` (the recommended Python tool installer).
8
+
9
+ If this project is published on PyPI:
10
+
11
+ ```sh
12
+ python -m pip install 'seraph-mcp[mcp]'
13
+ ```
14
+
15
+ If the user provides a Seraph source checkout instead:
16
+
17
+ ```sh
18
+ cd /path/to/seraph
19
+ uv tool install --from '.[mcp]' seraph-mcp
20
+ ```
21
+
22
+ Do not install Seraph into the target project or modify that project's dependencies. The user may choose another install method; preserve their existing environment and report any missing prerequisite clearly.
23
+
24
+ ## Register with the current agent
25
+
26
+ 1. Identify the current agent host and use its documented MCP setup mechanism (built-in setup command, supported CLI, settings UI, or documented config file).
27
+ 2. Register a **local stdio** MCP server named `seraph` with command `seraph-mcp` and no fixed repository environment variable.
28
+ 3. If `seraph-mcp` is not on the host's PATH, use the absolute executable path returned by `uv tool dir --bin` plus `/seraph-mcp` (macOS/Linux) or the equivalent path for the OS. Alternatively, register `uvx --from 'seraph-mcp[mcp]' seraph-mcp` when the host supports `uvx` and package downloads.
29
+ 4. Follow the host's instructions for reloading MCP servers, then confirm the Seraph tools are available. Never claim setup succeeded until the host confirms the server connected.
30
+
31
+ Do not assume a Codex, Claude, Cursor, VS Code, or other host config format. If the current host has no MCP support, explain that and give the user the relevant alternative instead of writing a guessed config.
32
+
33
+ ## Use it correctly
34
+
35
+ - Pass `repo` as the absolute path to the local Git checkout **on every tool call**.
36
+ - Use the active workspace path for questions about the current project.
37
+ - If the user asks about a different local project, use that project's path for that call; do not silently keep searching the active workspace.
38
+ - If the intended repository path is unclear, ask which local checkout they mean.
39
+ - Seraph does not upload repository contents. Its SQLite index is created under `<repo>/.seraph/` and is ignored by Git.
40
+ - Search tools index the requested Git version on first use. `index_repository` is available when the user wants indexing to happen explicitly first.
41
+ - The default setup is local BM25 and requires no API key.
42
+
43
+ ## Check setup
44
+
45
+ Call `index_repository` with the active workspace path, then call `search_code` with a question grounded in that repository and the same `repo` path. Report whether both calls succeeded and where the index was stored.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Prethiv Sriman D, Akshith Rajesh, Akhilan VTM, Sanjeev TS — Team floppydisk
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,115 @@
1
+ Metadata-Version: 2.5
2
+ Name: seraph-mcp
3
+ Version: 0.1.0
4
+ Summary: Query-adaptive, evolution-aware code retrieval
5
+ License: MIT License
6
+
7
+ Copyright (c) 2026 Prethiv Sriman D, Akshith Rajesh, Akhilan VTM, Sanjeev TS — Team floppydisk
8
+
9
+ Permission is hereby granted, free of charge, to any person obtaining a copy
10
+ of this software and associated documentation files (the "Software"), to deal
11
+ in the Software without restriction, including without limitation the rights
12
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
13
+ copies of the Software, and to permit persons to whom the Software is
14
+ furnished to do so, subject to the following conditions:
15
+
16
+ The above copyright notice and this permission notice shall be included in all
17
+ copies or substantial portions of the Software.
18
+
19
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
20
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
21
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
22
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
23
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
24
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
25
+ SOFTWARE.
26
+ License-File: LICENSE
27
+ Requires-Python: <3.13,>=3.12
28
+ Requires-Dist: bm25s>=0.2.6
29
+ Requires-Dist: numpy>=1.26
30
+ Requires-Dist: pydantic>=2.7
31
+ Requires-Dist: pystemmer>=2.2
32
+ Requires-Dist: pyyaml>=6.0
33
+ Requires-Dist: scikit-learn>=1.5
34
+ Requires-Dist: tree-sitter-go>=0.23
35
+ Requires-Dist: tree-sitter-javascript>=0.23
36
+ Requires-Dist: tree-sitter-python>=0.23
37
+ Requires-Dist: tree-sitter-typescript>=0.23
38
+ Requires-Dist: tree-sitter>=0.23
39
+ Provides-Extra: api
40
+ Requires-Dist: google-genai>=1.20; extra == 'api'
41
+ Requires-Dist: openai>=1.40; extra == 'api'
42
+ Provides-Extra: cpu
43
+ Requires-Dist: einops>=0.8; extra == 'cpu'
44
+ Requires-Dist: faiss-cpu>=1.8; extra == 'cpu'
45
+ Requires-Dist: sentence-transformers>=6.1; extra == 'cpu'
46
+ Requires-Dist: torch>=2.4; extra == 'cpu'
47
+ Requires-Dist: transformers>=5.19; extra == 'cpu'
48
+ Provides-Extra: gpu
49
+ Requires-Dist: einops>=0.8; extra == 'gpu'
50
+ Requires-Dist: faiss-cpu>=1.8; extra == 'gpu'
51
+ Requires-Dist: sentence-transformers>=6.1; extra == 'gpu'
52
+ Requires-Dist: torch>=2.4; extra == 'gpu'
53
+ Requires-Dist: transformers>=5.19; extra == 'gpu'
54
+ Provides-Extra: llm
55
+ Requires-Dist: openai>=1.40; extra == 'llm'
56
+ Provides-Extra: mcp
57
+ Requires-Dist: mcp<2,>=1.2; extra == 'mcp'
58
+ Description-Content-Type: text/markdown
59
+
60
+ # Seraph
61
+
62
+ Local code search and Git-aware code intelligence for coding agents. Seraph indexes a local Git checkout into a SQLite database and exposes search, symbol, dependency, history, and version comparison tools over MCP.
63
+
64
+ ## Install from a checkout
65
+
66
+ Requires Python 3.12 and Git. With `uv`:
67
+
68
+ ```sh
69
+ uv tool install --from '.[mcp]' seraph
70
+ ```
71
+
72
+ For development directly from this directory:
73
+
74
+ ```sh
75
+ uv sync --extra mcp
76
+ uv run seraph-mcp
77
+ ```
78
+
79
+ The published package install, once released to PyPI, is:
80
+
81
+ ```sh
82
+ python -m pip install 'seraph-mcp[mcp]'
83
+ ```
84
+
85
+ ## Agent setup
86
+
87
+ Give your coding agent [`AGENT_SETUP.md`](AGENT_SETUP.md). It explains how to install Seraph and register its local MCP server for that agent. MCP configuration differs by host, so the agent should use the host's own supported MCP setup mechanism rather than assume one universal config file.
88
+
89
+ Seraph runs on the user's machine. Each tool accepts `repo`, the absolute path to the local Git checkout to search. The agent should pass the active workspace path on every call; if asked about another checkout, pass that checkout's path. The index is stored at `<repo>/.seraph/index.sqlite` and ignored by Git.
90
+
91
+ ## MCP tools
92
+
93
+ - `search_code`: retrieve relevant code from the selected checkout.
94
+ - `search_at_version`: retrieve code from a commit, branch, or tag.
95
+ - `search_history`: search indexed Git versions.
96
+ - `find_symbol`: locate a function, class, or method.
97
+ - `find_dependencies`: walk call, import, and inheritance relationships.
98
+ - `compare_versions`: compare symbols, dependencies, and commits.
99
+ - `index_repository`: explicitly index a Git version; searches index automatically when needed.
100
+
101
+ The default retrieval setup is local and needs no API key. Benchmark/evaluation code and submission assets are intentionally excluded from this slim distribution.
102
+
103
+ ## CLI
104
+
105
+ ```sh
106
+ seraph --repo /path/to/repo index
107
+ seraph --repo /path/to/repo search "where is request authentication checked?"
108
+ seraph --repo /path/to/repo symbol authenticate_request
109
+ seraph --repo /path/to/repo deps authenticate_request
110
+ seraph --repo /path/to/repo compare <old-commit> <new-commit>
111
+ ```
112
+
113
+ ## License
114
+
115
+ See [`LICENSE`](LICENSE).
@@ -0,0 +1,56 @@
1
+ # Seraph
2
+
3
+ Local code search and Git-aware code intelligence for coding agents. Seraph indexes a local Git checkout into a SQLite database and exposes search, symbol, dependency, history, and version comparison tools over MCP.
4
+
5
+ ## Install from a checkout
6
+
7
+ Requires Python 3.12 and Git. With `uv`:
8
+
9
+ ```sh
10
+ uv tool install --from '.[mcp]' seraph
11
+ ```
12
+
13
+ For development directly from this directory:
14
+
15
+ ```sh
16
+ uv sync --extra mcp
17
+ uv run seraph-mcp
18
+ ```
19
+
20
+ The published package install, once released to PyPI, is:
21
+
22
+ ```sh
23
+ python -m pip install 'seraph-mcp[mcp]'
24
+ ```
25
+
26
+ ## Agent setup
27
+
28
+ Give your coding agent [`AGENT_SETUP.md`](AGENT_SETUP.md). It explains how to install Seraph and register its local MCP server for that agent. MCP configuration differs by host, so the agent should use the host's own supported MCP setup mechanism rather than assume one universal config file.
29
+
30
+ Seraph runs on the user's machine. Each tool accepts `repo`, the absolute path to the local Git checkout to search. The agent should pass the active workspace path on every call; if asked about another checkout, pass that checkout's path. The index is stored at `<repo>/.seraph/index.sqlite` and ignored by Git.
31
+
32
+ ## MCP tools
33
+
34
+ - `search_code`: retrieve relevant code from the selected checkout.
35
+ - `search_at_version`: retrieve code from a commit, branch, or tag.
36
+ - `search_history`: search indexed Git versions.
37
+ - `find_symbol`: locate a function, class, or method.
38
+ - `find_dependencies`: walk call, import, and inheritance relationships.
39
+ - `compare_versions`: compare symbols, dependencies, and commits.
40
+ - `index_repository`: explicitly index a Git version; searches index automatically when needed.
41
+
42
+ The default retrieval setup is local and needs no API key. Benchmark/evaluation code and submission assets are intentionally excluded from this slim distribution.
43
+
44
+ ## CLI
45
+
46
+ ```sh
47
+ seraph --repo /path/to/repo index
48
+ seraph --repo /path/to/repo search "where is request authentication checked?"
49
+ seraph --repo /path/to/repo symbol authenticate_request
50
+ seraph --repo /path/to/repo deps authenticate_request
51
+ seraph --repo /path/to/repo compare <old-commit> <new-commit>
52
+ ```
53
+
54
+ ## License
55
+
56
+ See [`LICENSE`](LICENSE).
@@ -0,0 +1,86 @@
1
+ [project]
2
+ name = "seraph-mcp"
3
+ version = "0.1.0"
4
+ description = "Query-adaptive, evolution-aware code retrieval"
5
+ readme = "README.md"
6
+ requires-python = ">=3.12,<3.13"
7
+ license = { file = "LICENSE" }
8
+ dependencies = [
9
+ "numpy>=1.26",
10
+ "pydantic>=2.7",
11
+ "pyyaml>=6.0",
12
+ "scikit-learn>=1.5",
13
+ "bm25s>=0.2.6",
14
+ "PyStemmer>=2.2",
15
+ "tree-sitter>=0.23",
16
+ "tree-sitter-python>=0.23",
17
+ "tree-sitter-typescript>=0.23",
18
+ "tree-sitter-javascript>=0.23",
19
+ "tree-sitter-go>=0.23",
20
+ ]
21
+
22
+ [project.optional-dependencies]
23
+ cpu = [
24
+ "torch>=2.4",
25
+ "transformers>=5.19",
26
+ "sentence-transformers>=6.1",
27
+ "faiss-cpu>=1.8",
28
+ "einops>=0.8",
29
+ ]
30
+ gpu = [
31
+ "torch>=2.4",
32
+ "transformers>=5.19",
33
+ "sentence-transformers>=6.1",
34
+ "faiss-cpu>=1.8",
35
+ "einops>=0.8",
36
+ ]
37
+ llm = ["openai>=1.40"]
38
+ api = [
39
+ "google-genai>=1.20",
40
+ "openai>=1.40",
41
+ ]
42
+ mcp = ["mcp>=1.2,<2"]
43
+
44
+ [project.scripts]
45
+ seraph = "seraph.cli:main"
46
+ seraph-mcp = "seraph.server:main"
47
+
48
+ [dependency-groups]
49
+ dev = ["pytest>=8.2", "ruff>=0.6"]
50
+
51
+ [build-system]
52
+ requires = ["hatchling"]
53
+ build-backend = "hatchling.build"
54
+
55
+ [tool.hatch.build.targets.wheel]
56
+ packages = ["src/seraph"]
57
+
58
+ [tool.uv]
59
+ python-preference = "managed"
60
+ conflicts = [[{ extra = "cpu" }, { extra = "gpu" }]]
61
+
62
+ [tool.uv.sources]
63
+ torch = [
64
+ { index = "pytorch-cpu", extra = "cpu" },
65
+ { index = "pytorch-cu126", extra = "gpu" },
66
+ ]
67
+
68
+ [[tool.uv.index]]
69
+ name = "pytorch-cpu"
70
+ url = "https://download.pytorch.org/whl/cpu"
71
+ explicit = true
72
+
73
+ [[tool.uv.index]]
74
+ name = "pytorch-cu126"
75
+ url = "https://download.pytorch.org/whl/cu126"
76
+ explicit = true
77
+
78
+ [tool.ruff]
79
+ line-length = 100
80
+ target-version = "py312"
81
+ src = ["src", "tests"]
82
+
83
+ [tool.ruff.lint]
84
+ select = ["E", "F", "I", "B", "UP", "SIM"]
85
+ ignore = ["E501", "B008"]
86
+
@@ -0,0 +1 @@
1
+ .[mcp]
@@ -0,0 +1,5 @@
1
+ """Seraph code intelligence engine."""
2
+
3
+ from .index import Chunk, IndexStats, SearchHit, VersionedIndex
4
+
5
+ __all__ = ["Chunk", "IndexStats", "SearchHit", "VersionedIndex"]
File without changes
@@ -0,0 +1,102 @@
1
+ """Shared plumbing for API backends: rate limiting, retries, persistent per-item cache."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import hashlib
6
+ import logging
7
+ import os
8
+ import re
9
+ import sqlite3
10
+ import threading
11
+ import time
12
+ from collections.abc import Callable
13
+ from pathlib import Path
14
+
15
+ log = logging.getLogger(__name__)
16
+
17
+ _RETRYABLE = ("429", "RESOURCE_EXHAUSTED", "503", "UNAVAILABLE", "500", "INTERNAL", "DEADLINE", "timed out")
18
+
19
+
20
+ class MissingAPIKeyError(RuntimeError):
21
+ pass
22
+
23
+
24
+ def require_key(env: str) -> str:
25
+ key = os.environ.get(env)
26
+ if not key:
27
+ raise MissingAPIKeyError(
28
+ f"{env} is not set. Set it (PowerShell: setx {env} <key>, then reopen the terminal) "
29
+ "or switch this component to a local model in the config."
30
+ )
31
+ return key
32
+
33
+
34
+ class RateLimiter:
35
+ """Spaces calls to at most `rpm` per minute."""
36
+
37
+ def __init__(self, rpm: float) -> None:
38
+ self.interval = 60.0 / rpm if rpm > 0 else 0.0
39
+ self._next = 0.0
40
+ self._lock = threading.Lock()
41
+
42
+ def wait(self) -> None:
43
+ with self._lock:
44
+ now = time.monotonic()
45
+ if now < self._next:
46
+ time.sleep(self._next - now)
47
+ self._next = max(now, self._next) + self.interval
48
+
49
+
50
+ class QuotaExhaustedError(RuntimeError):
51
+ """A daily quota is used up; retrying today will not help. Switch keys or wait."""
52
+
53
+
54
+ _RETRY_DELAY = re.compile(r"retry(?:Delay)?['\"]?\s*[:=]?\s*['\"]?(\d+(?:\.\d+)?)s", re.I)
55
+
56
+
57
+ def with_retries[T](fn: Callable[[], T], attempts: int = 8, base_delay: float = 4.0) -> T:
58
+ for i in range(attempts):
59
+ try:
60
+ return fn()
61
+ except Exception as e:
62
+ msg = str(e)
63
+ if "PerDay" in msg:
64
+ raise QuotaExhaustedError(f"daily API quota exhausted: {msg[:300]}") from e
65
+ if i == attempts - 1 or not any(code in msg for code in _RETRYABLE):
66
+ raise
67
+ hint = _RETRY_DELAY.search(msg)
68
+ delay = float(hint.group(1)) + 1 if hint else min(base_delay * 2**i, 120.0)
69
+ log.warning("API call failed (%s); retrying in %.0fs", msg[:120], delay)
70
+ time.sleep(delay)
71
+ raise AssertionError("unreachable")
72
+
73
+
74
+ class KVCache:
75
+ """SQLite key -> bytes store; survives interrupted runs (free-tier daily limits)."""
76
+
77
+ def __init__(self, path: str | Path) -> None:
78
+ path = Path(path)
79
+ path.parent.mkdir(parents=True, exist_ok=True)
80
+ self._db = sqlite3.connect(path, check_same_thread=False)
81
+ self._db.execute("CREATE TABLE IF NOT EXISTS kv (k TEXT PRIMARY KEY, v BLOB)")
82
+ self._db.commit()
83
+
84
+ @staticmethod
85
+ def key(*parts: str) -> str:
86
+ h = hashlib.sha1()
87
+ for p in parts:
88
+ h.update(p.encode())
89
+ h.update(b"\x00")
90
+ return h.hexdigest()
91
+
92
+ def get_many(self, keys: list[str]) -> dict[str, bytes]:
93
+ out: dict[str, bytes] = {}
94
+ for i in range(0, len(keys), 500):
95
+ chunk = keys[i : i + 500]
96
+ q = f"SELECT k, v FROM kv WHERE k IN ({','.join('?' * len(chunk))})"
97
+ out.update(self._db.execute(q, chunk).fetchall())
98
+ return out
99
+
100
+ def put_many(self, items: dict[str, bytes]) -> None:
101
+ self._db.executemany("INSERT OR REPLACE INTO kv VALUES (?, ?)", items.items())
102
+ self._db.commit()
@@ -0,0 +1,186 @@
1
+ """Gemini API backends: embeddings and an LLM listwise reranker."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ import logging
7
+ import re
8
+ from collections.abc import Sequence
9
+ from pathlib import Path
10
+
11
+ import numpy as np
12
+
13
+ from seraph.backends.common import KVCache, RateLimiter, require_key, with_retries
14
+ from seraph.config import ModelSpec
15
+ from seraph.retrieval.reranker import head_tail
16
+ from seraph.retrieval.vectors import normalize
17
+ from seraph.types import ScoredChunk
18
+
19
+ log = logging.getLogger(__name__)
20
+
21
+
22
+ def make_client(api_key_env: str = "GEMINI_API_KEY", client=None):
23
+ if client is not None:
24
+ return client
25
+ from google import genai
26
+
27
+ return genai.Client(api_key=require_key(api_key_env))
28
+
29
+
30
+ class GeminiEmbedder:
31
+ """`gemini-embedding-001` with task types; every text is cached individually."""
32
+
33
+ def __init__(self, spec: ModelSpec, cache_dir: str | Path, client=None) -> None:
34
+ o = spec.options
35
+ self.spec = spec
36
+ self.name = spec.name
37
+ self.model = spec.hf_id
38
+ self.dim = int(o.get("output_dim", 768))
39
+ self.revision = f"{self.model}-d{self.dim}"
40
+ self.query_task = o.get("query_task", "CODE_RETRIEVAL_QUERY")
41
+ self.doc_task = o.get("doc_task", "RETRIEVAL_DOCUMENT")
42
+ self.batch = int(o.get("batch", 100))
43
+ self.max_chars = int(o.get("max_chars", 6000))
44
+ self._client = client
45
+ self._limiter = RateLimiter(float(o.get("rpm", 90)))
46
+ self._cache = KVCache(Path(cache_dir) / "api" / f"{spec.name}.sqlite")
47
+
48
+ @property
49
+ def client(self):
50
+ self._client = make_client(self.spec.options.get("api_key_env", "GEMINI_API_KEY"), self._client)
51
+ return self._client
52
+
53
+ def _call(self, texts: list[str], task: str) -> np.ndarray:
54
+ from google.genai import types
55
+
56
+ def go():
57
+ self._limiter.wait()
58
+ resp = self.client.models.embed_content(
59
+ model=self.model,
60
+ contents=texts,
61
+ config=types.EmbedContentConfig(task_type=task, output_dimensionality=self.dim),
62
+ )
63
+ return np.array([e.values for e in resp.embeddings], dtype=np.float32)
64
+
65
+ return with_retries(go)
66
+
67
+ def _encode(self, texts: Sequence[str], task: str, prefix: str) -> np.ndarray:
68
+ texts = [(prefix + t)[: self.max_chars] or " " for t in texts]
69
+ keys = [KVCache.key(self.revision, task, t) for t in texts]
70
+ have = self._cache.get_many(list(set(keys)))
71
+ missing = sorted({k: t for k, t in zip(keys, texts, strict=True) if k not in have}.items())
72
+ for i in range(0, len(missing), self.batch):
73
+ part = missing[i : i + self.batch]
74
+ vecs = self._call([t for _, t in part], task)
75
+ new = {k: v.tobytes() for (k, _), v in zip(part, vecs, strict=True)}
76
+ self._cache.put_many(new)
77
+ have.update(new)
78
+ if len(missing) > self.batch:
79
+ log.info("embedded %d/%d", min(i + self.batch, len(missing)), len(missing))
80
+ return normalize(np.stack([np.frombuffer(have[k], dtype=np.float32) for k in keys]))
81
+
82
+ def encode_queries(self, texts: Sequence[str]) -> np.ndarray:
83
+ return self._encode(texts, self.query_task, self.spec.query_prefix)
84
+
85
+ def encode_documents(self, texts: Sequence[str]) -> np.ndarray:
86
+ return self._encode(texts, self.doc_task, self.spec.doc_prefix)
87
+
88
+
89
+ LISTWISE_PROMPT = """You rank candidate code snippets by how likely each one solves the programming task.
90
+
91
+ Task:
92
+ {query}
93
+
94
+ Candidates:
95
+ {candidates}
96
+
97
+ Return ONLY a JSON array of the candidate numbers, best first, including every candidate once."""
98
+
99
+
100
+ class GeminiListwiseReranker:
101
+ """One Gemini call per query reorders the top-N candidates; results are cached per query+set."""
102
+
103
+ def __init__(self, spec: ModelSpec, cache_dir: str | Path, blend: float = 0.0, client=None) -> None:
104
+ o = spec.options
105
+ self.spec = spec
106
+ self.name = spec.name
107
+ self.model = spec.hf_id
108
+ self.revision = self.model
109
+ self.blend = blend
110
+ self.window = int(o.get("window", 20))
111
+ self.doc_chars = int(o.get("doc_chars", 1200))
112
+ self.query_chars = int(o.get("query_chars", 3000))
113
+ self._client = client
114
+ self._limiter = RateLimiter(float(o.get("rpm", 14)))
115
+ self._cache = KVCache(Path(cache_dir) / "api" / f"{spec.name}.sqlite")
116
+
117
+ @property
118
+ def client(self):
119
+ self._client = make_client(self.spec.options.get("api_key_env", "GEMINI_API_KEY"), self._client)
120
+ return self._client
121
+
122
+ def _prompt(self, query: str, cands: Sequence[ScoredChunk]) -> str:
123
+ body = "\n\n".join(
124
+ f"[{i + 1}]\n```\n{head_tail(c.chunk.text, self.doc_chars)}\n```" for i, c in enumerate(cands)
125
+ )
126
+ return LISTWISE_PROMPT.format(query=head_tail(query, self.query_chars), candidates=body)
127
+
128
+ def _order(self, query: str, cands: Sequence[ScoredChunk]) -> list[int]:
129
+ key = KVCache.key(self.model, query, *[c.chunk_hash for c in cands])
130
+ hit = self._cache.get_many([key])
131
+ if key in hit:
132
+ return json.loads(hit[key])
133
+ from google.genai import types
134
+
135
+ def go():
136
+ self._limiter.wait()
137
+ return self.client.models.generate_content(
138
+ model=self.model,
139
+ contents=self._prompt(query, cands),
140
+ config=types.GenerateContentConfig(temperature=0.0, response_mime_type="application/json"),
141
+ ).text
142
+
143
+ order = parse_order(with_retries(go) or "", len(cands))
144
+ self._cache.put_many({key: json.dumps(order).encode()})
145
+ return order
146
+
147
+ def rerank_batch(
148
+ self, queries: Sequence[str], candidates: Sequence[Sequence[ScoredChunk]], top_k: int
149
+ ) -> list[list[ScoredChunk]]:
150
+ out = []
151
+ for q, cands in zip(queries, candidates, strict=True):
152
+ head, rest = list(cands[: self.window]), list(cands[self.window :])
153
+ order = self._order(q, head) if head else []
154
+ n = len(head)
155
+ ranked = []
156
+ for rank, idx in enumerate(order):
157
+ c = head[idx]
158
+ s = 1.0 - rank / max(n, 1)
159
+ final = s if self.blend == 0 else (1 - self.blend) * s + self.blend * c.score
160
+ ranked.append(ScoredChunk(c.chunk, final, {**c.scores, "rerank": s}))
161
+ ranked.sort(key=lambda x: -x.score)
162
+ floor = min((r.score for r in ranked), default=0.0)
163
+ ranked += [ScoredChunk(c.chunk, floor - (j + 1) * 1e-3, c.scores) for j, c in enumerate(rest)]
164
+ out.append(ranked[:top_k])
165
+ return out
166
+
167
+ def rerank(self, query: str, candidates: Sequence[ScoredChunk], top_k: int) -> list[ScoredChunk]:
168
+ return self.rerank_batch([query], [candidates], top_k)[0]
169
+
170
+
171
+ def parse_order(text: str, n: int) -> list[int]:
172
+ """Parse a 1-based ranking; drop invalid/duplicate ids and append any the model skipped."""
173
+ try:
174
+ raw = json.loads(text)
175
+ nums = raw if isinstance(raw, list) else raw.get("ranking", [])
176
+ except (json.JSONDecodeError, AttributeError):
177
+ nums = re.findall(r"\d+", text)
178
+ seen: list[int] = []
179
+ for x in nums:
180
+ try:
181
+ i = int(x) - 1
182
+ except (TypeError, ValueError):
183
+ continue
184
+ if 0 <= i < n and i not in seen:
185
+ seen.append(i)
186
+ return seen + [i for i in range(n) if i not in seen]