tools: generate the hanzo surface from cloud's OpenAPI registry; delete the vector fake
The `hanzo` dispatcher hand-listed 10 services in SERVICE_TOOL_PATHS, so the other ~150 cloud products were unreachable until someone cut a Python release. Project the surface from cloud's /v1/openapi.json instead: cloud tags every operation with its product (the first path segment after /v1/), so that document already IS a service/action catalog. 166 services are now reachable, and a newly mounted app is callable the moment cloud serves it. Still exactly one MCP tool — agents degrade badly with hundreds — now shaped hanzo(service, action, params, method), with `services` listing the catalog and an empty action listing a service's actions. Delete the local Infinity vector store (3,546 lines). It shipped no embedder, so its only possible output was random vectors: it ranked "Bananas are yellow" above an auth document for the query "authentication oauth jwt". A previous fix put the mock behind HANZO_VECTOR_ALLOW_MOCK, but a fake that a flag can re-enable is still a fake, and a tool that lies to an agent is worse than a missing one. The package keeps only the cloud-backed VectorTool, and a test now fails if any module in it imports `random`. Make the legacy-tool gate honest: hide the per-service tools only once the unified tool genuinely imports. It does not import against released hanzo-tools (HanzoCloud is newer), and disabling them unconditionally would have left no cloud tools at all. - spec.py: catalog projection, cache keyed per spec source, stale cache beats failing a call. HANZO_OPENAPI_URL decouples catalog from target so a partial local host can be driven with the full registry. - HanzoCloud.call: one generic seam for the verbs the spec names (PUT/PATCH/ DELETE); auth also resolves the hanzo CLI's IAM session, memoized. - 40 new tests covering action naming, template binding, method inference, param split, and cache fallback. Co-authored-by: Hanzo Dev <dev@hanzo.ai>
This commit is contained in:
zeekay
@@ -21,6 +21,35 @@ if TYPE_CHECKING:
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Per-service tools the unified `hanzo` tool replaces, each now reachable as
|
||||
# hanzo(service=...) straight from cloud's OpenAPI registry.
|
||||
SUPERSEDED_BY_HANZO = (
|
||||
"api",
|
||||
"auth",
|
||||
"billing",
|
||||
"commerce",
|
||||
"iam",
|
||||
"ingress",
|
||||
"kms",
|
||||
"mpc",
|
||||
"paas",
|
||||
"team",
|
||||
)
|
||||
|
||||
|
||||
def _unified_hanzo_available() -> bool:
|
||||
"""Whether the unified `hanzo` tool can actually be constructed.
|
||||
|
||||
Importing is the honest test: it is what the loader will do, so a success
|
||||
here means the replacement really is there to take over.
|
||||
"""
|
||||
try:
|
||||
from hanzo_tools.api.hanzo_tool import HanzoTool # noqa: F401
|
||||
except Exception as exc: # ImportError, or a broken dependency
|
||||
logger.warning(f"unified hanzo tool not importable: {exc}")
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def register_all_tools(
|
||||
mcp_server: "FastMCP",
|
||||
@@ -164,22 +193,19 @@ def register_all_tools(
|
||||
# MCP tools
|
||||
resolved_enabled_tools["mcp"] = is_tool_enabled("mcp", True)
|
||||
|
||||
# Unified Hanzo platform surface.
|
||||
# No backwards compatibility: expose only `hanzo` and hide legacy per-service tools.
|
||||
# Unified Hanzo platform surface: `hanzo` covers every cloud service by
|
||||
# projecting them from the OpenAPI registry, so the per-service tools are
|
||||
# redundant. Hide them only once the replacement genuinely exists —
|
||||
# unconditionally disabling them would strand the user with no cloud tools
|
||||
# at all on any import error in the unified surface.
|
||||
resolved_enabled_tools["hanzo"] = is_tool_enabled("hanzo", True)
|
||||
for legacy_tool in [
|
||||
"api",
|
||||
"auth",
|
||||
"billing",
|
||||
"commerce",
|
||||
"iam",
|
||||
"ingress",
|
||||
"kms",
|
||||
"mpc",
|
||||
"paas",
|
||||
"team",
|
||||
]:
|
||||
resolved_enabled_tools[legacy_tool] = False
|
||||
if resolved_enabled_tools["hanzo"] and _unified_hanzo_available():
|
||||
for legacy_tool in SUPERSEDED_BY_HANZO:
|
||||
resolved_enabled_tools[legacy_tool] = False
|
||||
else:
|
||||
logger.warning(
|
||||
"unified `hanzo` tool unavailable — keeping per-service tools enabled"
|
||||
)
|
||||
|
||||
# Jupyter tools
|
||||
for tool in PACKAGE_TOOL_PREFIXES.get("jupyter", []):
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hanzo-mcp"
|
||||
version = "0.15.12"
|
||||
version = "0.15.13"
|
||||
description = "The Zen of Hanzo MCP: One server to rule them all. The ultimate MCP that orchestrates all others."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
|
||||
@@ -1,14 +1,16 @@
|
||||
"""Unified Hanzo platform tool.
|
||||
"""Unified Hanzo platform tool, projected from cloud's OpenAPI registry.
|
||||
|
||||
Provides a compact `hanzo` surface that routes to Hanzo service tools
|
||||
(`auth`, `billing`, `commerce`, `iam`, `ingress`, `kms`, `mpc`, `paas`,
|
||||
`team`, and generic `api`).
|
||||
One MCP tool reaches every service cloud serves. The service/action surface is
|
||||
derived from ``/v1/openapi.json`` at runtime rather than hand-listed here, so a
|
||||
newly mounted app is callable the moment cloud serves it — no Python release,
|
||||
and no list that silently drifts behind the platform.
|
||||
|
||||
One tool, not one per service: agents degrade badly with hundreds of tools, and
|
||||
the spec already carries the routing information a dispatcher needs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import inspect
|
||||
import json
|
||||
from typing import Annotated, Any, final, override
|
||||
|
||||
@@ -16,168 +18,154 @@ from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
from pydantic import Field
|
||||
|
||||
from hanzo_tools.core import BaseTool, auto_timeout, create_tool_context
|
||||
from hanzo_tools.core import BaseTool, HanzoCloud, auto_timeout, create_tool_context
|
||||
|
||||
SERVICE_TOOL_PATHS: dict[str, str] = {
|
||||
"api": "hanzo_tools.api.api_tool:APITool",
|
||||
"auth": "hanzo_tools.auth.login_tool:LoginTool",
|
||||
"billing": "hanzo_tools.billing.billing_tool:BillingTool",
|
||||
"commerce": "hanzo_tools.commerce.commerce_tool:CommerceTool",
|
||||
"iam": "hanzo_tools.iam.iam_tool:IAMTool",
|
||||
"ingress": "hanzo_tools.ingress.ingress_tool:IngressTool",
|
||||
"kms": "hanzo_tools.kms.kms_tool:KMSTool",
|
||||
"mpc": "hanzo_tools.mpc.mpc_tool:MPCTool",
|
||||
"paas": "hanzo_tools.paas.paas_tool:PaaSTool",
|
||||
"team": "hanzo_tools.team.team_tool:TeamTool",
|
||||
}
|
||||
from . import spec as openapi
|
||||
|
||||
# Names that address the catalog itself rather than a cloud service.
|
||||
CATALOG_SERVICES = frozenset({"services", "list"})
|
||||
REFRESH_SERVICES = frozenset({"refresh", "reload"})
|
||||
|
||||
# Convenience synonyms for services whose product tag is not the obvious word.
|
||||
SERVICE_ALIASES: dict[str, str] = {
|
||||
"platform": "paas",
|
||||
"identity": "iam",
|
||||
"payments": "billing",
|
||||
"store": "commerce",
|
||||
"knowledge": "kb",
|
||||
"platform": "paas",
|
||||
"rag": "kb",
|
||||
}
|
||||
|
||||
|
||||
@final
|
||||
class HanzoTool(BaseTool):
|
||||
"""Unified tool for Hanzo platform services."""
|
||||
"""Unified tool for every Hanzo cloud service."""
|
||||
|
||||
name = "hanzo"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._delegates: dict[str, BaseTool] = {}
|
||||
self._catalog: openapi.Catalog | None = None
|
||||
self._source: str = ""
|
||||
self._cloud: HanzoCloud | None = None
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
return """Unified Hanzo platform tool.
|
||||
return """Unified Hanzo platform tool — every service cloud serves.
|
||||
|
||||
Use one `hanzo` tool surface for service operations across:
|
||||
- auth
|
||||
- billing
|
||||
- commerce
|
||||
- iam
|
||||
- ingress
|
||||
- kms
|
||||
- mpc
|
||||
- paas
|
||||
- team
|
||||
- api (generic OpenAPI bridge)
|
||||
The service/action surface is generated from cloud's live OpenAPI registry, so
|
||||
it always matches what the platform actually serves.
|
||||
|
||||
Parameters:
|
||||
- service: Target Hanzo service
|
||||
- action: Service-specific action
|
||||
- args: JSON object string for service-specific parameters
|
||||
- service: Product tag (the first path segment after /v1/). "services" lists them.
|
||||
- action: Path within the service ("plans" -> /v1/billing/plans). Omit to list
|
||||
a service's actions.
|
||||
- params: JSON object string; sent as the body for POST/PUT/PATCH, else as the
|
||||
query string. Values also fill {templated} path segments.
|
||||
- method: Override the HTTP method when an action serves several.
|
||||
|
||||
Examples:
|
||||
hanzo(service="auth", action="status")
|
||||
hanzo(service="commerce", action="orders")
|
||||
hanzo(service="iam", action="users", args='{"owner":"hanzo"}')
|
||||
hanzo(service="api", action="list")
|
||||
hanzo(service="services")
|
||||
hanzo(service="billing")
|
||||
hanzo(service="billing", action="plans")
|
||||
hanzo(service="kb", action="search", params='{"query":"authentication oauth jwt"}')
|
||||
hanzo(service="refresh")
|
||||
"""
|
||||
|
||||
def _normalize_service(self, service: str) -> str:
|
||||
key = (service or "").strip().lower().replace("-", "_")
|
||||
key = SERVICE_ALIASES.get(key, key)
|
||||
return key
|
||||
def _get_catalog(self, refresh: bool = False) -> openapi.Catalog:
|
||||
if self._catalog is None or refresh:
|
||||
document, source = openapi.load(refresh=refresh)
|
||||
self._catalog = openapi.Catalog(document)
|
||||
self._source = source
|
||||
return self._catalog
|
||||
|
||||
def _load_delegate(self, service: str) -> BaseTool:
|
||||
if service in self._delegates:
|
||||
return self._delegates[service]
|
||||
def _get_cloud(self) -> HanzoCloud:
|
||||
if self._cloud is None:
|
||||
self._cloud = HanzoCloud()
|
||||
return self._cloud
|
||||
|
||||
path = SERVICE_TOOL_PATHS.get(service)
|
||||
if not path:
|
||||
available = ", ".join(sorted(SERVICE_TOOL_PATHS.keys()))
|
||||
raise ValueError(
|
||||
f"Unknown service '{service}'. Available services: {available}"
|
||||
)
|
||||
@staticmethod
|
||||
def _normalize(service: str) -> str:
|
||||
key = (service or "").strip().lower()
|
||||
return SERVICE_ALIASES.get(key, key)
|
||||
|
||||
module_name, class_name = path.split(":")
|
||||
module = importlib.import_module(module_name)
|
||||
cls = getattr(module, class_name)
|
||||
tool = cls()
|
||||
self._delegates[service] = tool
|
||||
return tool
|
||||
|
||||
def _parse_args(self, args: str | None) -> dict[str, Any]:
|
||||
if not args:
|
||||
@staticmethod
|
||||
def _parse_params(params: str | None) -> dict[str, Any]:
|
||||
if not params:
|
||||
return {}
|
||||
try:
|
||||
parsed = json.loads(args)
|
||||
parsed = json.loads(params)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"args must be valid JSON object string: {exc}") from exc
|
||||
|
||||
raise ValueError(f"params must be a JSON object string: {exc}") from exc
|
||||
if not isinstance(parsed, dict):
|
||||
raise ValueError("args must decode to a JSON object")
|
||||
raise ValueError("params must decode to a JSON object")
|
||||
return parsed
|
||||
|
||||
async def _delegate_call(
|
||||
self,
|
||||
tool: BaseTool,
|
||||
ctx: MCPContext,
|
||||
payload: dict[str, Any],
|
||||
) -> str:
|
||||
sig = inspect.signature(tool.call)
|
||||
params = sig.parameters
|
||||
accepts_kwargs = any(
|
||||
p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()
|
||||
)
|
||||
|
||||
if accepts_kwargs:
|
||||
return await tool.call(ctx, **payload)
|
||||
|
||||
allowed = {
|
||||
name
|
||||
for name, param in params.items()
|
||||
if name not in {"self", "ctx"}
|
||||
and param.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)
|
||||
}
|
||||
filtered = {k: v for k, v in payload.items() if k in allowed}
|
||||
return await tool.call(ctx, **filtered)
|
||||
|
||||
@override
|
||||
@auto_timeout("hanzo")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
service: str = "api",
|
||||
action: str = "list",
|
||||
args: str | None = None,
|
||||
service: str = "services",
|
||||
action: str = "",
|
||||
params: str | dict[str, Any] | None = None,
|
||||
method: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
service_key = self._normalize_service(service)
|
||||
if service_key in {"services", "list"}:
|
||||
return json.dumps(
|
||||
{
|
||||
"services": sorted(SERVICE_TOOL_PATHS.keys()),
|
||||
"aliases": SERVICE_ALIASES,
|
||||
"usage": 'hanzo(service="commerce", action="orders", args="{\\"query\\":\\"...\\\"}")',
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
|
||||
key = self._normalize(service)
|
||||
try:
|
||||
delegate = self._load_delegate(service_key)
|
||||
payload = {"action": action}
|
||||
payload.update(self._parse_args(args))
|
||||
payload.update({k: v for k, v in kwargs.items() if v is not None})
|
||||
return await self._delegate_call(delegate, ctx, payload)
|
||||
except Exception as exc:
|
||||
return json.dumps(
|
||||
{
|
||||
"error": str(exc),
|
||||
"service": service_key,
|
||||
"available_services": sorted(SERVICE_TOOL_PATHS.keys()),
|
||||
},
|
||||
indent=2,
|
||||
if key in REFRESH_SERVICES:
|
||||
catalog = self._get_catalog(refresh=True)
|
||||
return self._dump(
|
||||
{
|
||||
"refreshed": True,
|
||||
"source": self._source,
|
||||
"services": len(catalog.services),
|
||||
}
|
||||
)
|
||||
|
||||
catalog = self._get_catalog()
|
||||
|
||||
if key in CATALOG_SERVICES:
|
||||
summary = catalog.summary()
|
||||
return self._dump(
|
||||
{
|
||||
"source": self._source,
|
||||
"count": len(summary),
|
||||
"services": summary,
|
||||
"usage": 'hanzo(service="billing", action="plans")',
|
||||
}
|
||||
)
|
||||
|
||||
supplied = params if isinstance(params, dict) else self._parse_params(params)
|
||||
supplied = {**supplied, **{k: v for k, v in kwargs.items() if v is not None}}
|
||||
|
||||
if not action:
|
||||
return self._dump(catalog.describe(key))
|
||||
|
||||
route = catalog.resolve(key, action, method, has_params=bool(supplied))
|
||||
query, body = openapi.split_params(route, supplied)
|
||||
data = await self._get_cloud().call(route.method, route.path, query, body)
|
||||
return self._dump(
|
||||
{"service": key, "request": f"{route.method} {route.path}", "data": data}
|
||||
)
|
||||
except Exception as exc:
|
||||
return self._dump(self._error(key, exc))
|
||||
|
||||
def _error(self, service: str, exc: Exception) -> dict[str, Any]:
|
||||
out: dict[str, Any] = {"error": str(exc), "service": service}
|
||||
if self._catalog is not None and service not in self._catalog.services:
|
||||
out["hint"] = 'call hanzo(service="services") to list services'
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def _dump(payload: dict[str, Any]) -> str:
|
||||
return json.dumps(payload, indent=2, default=str)
|
||||
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register unified hanzo tool with explicit compact params."""
|
||||
"""Register the single unified hanzo tool."""
|
||||
tool_instance = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
@@ -186,24 +174,35 @@ Examples:
|
||||
str,
|
||||
Field(
|
||||
description=(
|
||||
"Target service: api, auth, billing, commerce, iam, ingress, "
|
||||
"kms, mpc, paas, team. Use 'services' to list."
|
||||
"Product tag, e.g. billing, iam, kb, paas, git. "
|
||||
'Use "services" to list them, "refresh" to refetch the spec.'
|
||||
)
|
||||
),
|
||||
] = "api",
|
||||
] = "services",
|
||||
action: Annotated[
|
||||
str,
|
||||
Field(description="Service action to execute (service-specific)."),
|
||||
] = "list",
|
||||
args: Annotated[
|
||||
Field(
|
||||
description=(
|
||||
"Path within the service, e.g. 'plans' for "
|
||||
"/v1/billing/plans. Omit to list the service's actions."
|
||||
)
|
||||
),
|
||||
] = "",
|
||||
params: Annotated[
|
||||
str | None,
|
||||
Field(
|
||||
description=(
|
||||
"JSON object string with service-specific parameters. "
|
||||
'Example: "{\\"query\\":\\"foo\\",\\"owner\\":\\"hanzo\\"}"'
|
||||
"JSON object string. Body for POST/PUT/PATCH, else query "
|
||||
'string. Example: "{\\"query\\":\\"oauth jwt\\"}"'
|
||||
)
|
||||
),
|
||||
] = None,
|
||||
method: Annotated[
|
||||
str | None,
|
||||
Field(description="Override HTTP method (GET/POST/PUT/PATCH/DELETE)."),
|
||||
] = None,
|
||||
ctx: MCPContext = None, # type: ignore[assignment]
|
||||
) -> str:
|
||||
return await tool_instance.call(ctx, service=service, action=action, args=args)
|
||||
return await tool_instance.call(
|
||||
ctx, service=service, action=action, params=params, method=method
|
||||
)
|
||||
|
||||
@@ -0,0 +1,290 @@
|
||||
"""Service catalog projected from cloud's live OpenAPI registry.
|
||||
|
||||
Cloud generates ``/v1/openapi.json`` from the router it actually serves, and
|
||||
tags every operation with its product — the first path segment after ``/v1/``.
|
||||
That document therefore already IS a service/action catalog, so we project the
|
||||
MCP surface from it instead of hand-listing services in Python: a new app is
|
||||
reachable the moment cloud serves it, with no Python release. Hand-listing is
|
||||
precisely what lets a tool surface drift out of date with the platform.
|
||||
|
||||
The catalog is a pure function of the spec. Fetching/caching is separate (see
|
||||
``load``) so the mapping can be tested without a network.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from collections.abc import Iterable
|
||||
from pathlib import Path
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
DEFAULT_BASE = "https://api.hanzo.ai"
|
||||
SPEC_PATH = "/v1/openapi.json"
|
||||
CACHE_TTL = 24 * 3600
|
||||
METHODS = ("get", "post", "put", "patch", "delete", "head", "options")
|
||||
# Methods that carry input as a JSON body; everything else takes a query string.
|
||||
BODY_METHODS = frozenset({"post", "put", "patch"})
|
||||
|
||||
|
||||
def api_base() -> str:
|
||||
"""Resolve the cloud base URL (no trailing slash)."""
|
||||
base = (
|
||||
os.environ.get("HANZO_API_BASE")
|
||||
or os.environ.get("HANZO_BASE_URL")
|
||||
or DEFAULT_BASE
|
||||
)
|
||||
return base.rstrip("/")
|
||||
|
||||
|
||||
class Route(NamedTuple):
|
||||
"""One concrete call: a method, a path with params already bound."""
|
||||
|
||||
method: str
|
||||
path: str
|
||||
# Params consumed by the path template, so the caller does not resend them.
|
||||
bound: dict[str, str]
|
||||
|
||||
|
||||
class Operation(NamedTuple):
|
||||
method: str
|
||||
path: str
|
||||
action: str
|
||||
typed: bool
|
||||
|
||||
|
||||
def _segments(path: str) -> list[str]:
|
||||
return [s for s in path.strip("/").split("/") if s]
|
||||
|
||||
|
||||
def action_of(path: str, tag: str) -> str:
|
||||
"""The action name for a path within its service.
|
||||
|
||||
``/v1/billing/plans`` under tag ``billing`` is the action ``plans``. The
|
||||
version and the service segment are dropped because they are already
|
||||
implied by the service, leaving the shortest thing a caller can name.
|
||||
"""
|
||||
segs = _segments(path)
|
||||
if segs and segs[0] == "v1":
|
||||
segs = segs[1:]
|
||||
if segs and segs[0] == tag:
|
||||
segs = segs[1:]
|
||||
return "/".join(segs)
|
||||
|
||||
|
||||
def _bind(want: str, have: str) -> dict[str, str] | None:
|
||||
"""Match a requested action against a (possibly templated) action.
|
||||
|
||||
Returns the path-param bindings, or None if it does not match. A literal
|
||||
action wins over a template because an exact match binds nothing.
|
||||
"""
|
||||
w, h = _segments(want), _segments(have)
|
||||
if len(w) != len(h):
|
||||
return None
|
||||
bound: dict[str, str] = {}
|
||||
for got, pat in zip(w, h, strict=True): # lengths checked above
|
||||
if len(pat) > 2 and pat[0] == "{" and pat[-1] == "}":
|
||||
bound[pat[1:-1]] = got
|
||||
elif got != pat:
|
||||
return None
|
||||
return bound
|
||||
|
||||
|
||||
class Catalog:
|
||||
"""Services and actions derived from an OpenAPI document."""
|
||||
|
||||
def __init__(self, spec: dict[str, Any]):
|
||||
self.spec = spec
|
||||
self._ops: dict[str, list[Operation]] = {}
|
||||
for path, item in (spec.get("paths") or {}).items():
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
for method, op in item.items():
|
||||
if method not in METHODS or not isinstance(op, dict):
|
||||
continue
|
||||
# An op is typed when cloud folded a real schema in; untyped ops
|
||||
# are router-shape only and take params on trust.
|
||||
typed = bool(op.get("requestBody") or op.get("responses"))
|
||||
for tag in op.get("tags") or ["_untagged"]:
|
||||
self._ops.setdefault(tag, []).append(
|
||||
Operation(method, path, action_of(path, tag), typed)
|
||||
)
|
||||
|
||||
@property
|
||||
def services(self) -> list[str]:
|
||||
return sorted(self._ops)
|
||||
|
||||
def operations(self, service: str) -> list[Operation]:
|
||||
return sorted(self._ops.get(service, []), key=lambda o: (o.action, o.method))
|
||||
|
||||
def summary(self) -> dict[str, Any]:
|
||||
"""Compact catalog: every service with its operation and typed counts."""
|
||||
return {
|
||||
svc: {
|
||||
"actions": len({o.action for o in ops}),
|
||||
"operations": len(ops),
|
||||
"typed": sum(1 for o in ops if o.typed),
|
||||
}
|
||||
for svc, ops in sorted(self._ops.items())
|
||||
}
|
||||
|
||||
def describe(self, service: str) -> dict[str, Any]:
|
||||
ops = self.operations(service)
|
||||
if not ops:
|
||||
raise KeyError(service)
|
||||
actions: dict[str, dict[str, Any]] = {}
|
||||
for o in ops:
|
||||
entry = actions.setdefault(
|
||||
o.action, {"methods": [], "path": o.path, "typed": False}
|
||||
)
|
||||
entry["methods"].append(o.method.upper())
|
||||
entry["typed"] = entry["typed"] or o.typed
|
||||
return {"service": service, "actions": actions}
|
||||
|
||||
def resolve(
|
||||
self,
|
||||
service: str,
|
||||
action: str,
|
||||
method: str | None = None,
|
||||
has_params: bool = False,
|
||||
) -> Route:
|
||||
"""Pick the one route named by (service, action).
|
||||
|
||||
Exact action matches beat templated ones. When several methods share an
|
||||
action and the caller did not name one, params imply a write (POST) and
|
||||
their absence implies a read (GET) — the intent the caller expressed.
|
||||
"""
|
||||
ops = self._ops.get(service)
|
||||
if not ops:
|
||||
raise KeyError(f"unknown service '{service}'")
|
||||
|
||||
want = (action or "").strip("/")
|
||||
exact: list[tuple[Operation, dict[str, str]]] = []
|
||||
templated: list[tuple[Operation, dict[str, str]]] = []
|
||||
for o in ops:
|
||||
bound = _bind(want, o.action)
|
||||
if bound is None:
|
||||
continue
|
||||
(exact if not bound else templated).append((o, bound))
|
||||
|
||||
candidates = exact or templated
|
||||
if not candidates:
|
||||
known = sorted({o.action for o in ops})[:20]
|
||||
raise KeyError(
|
||||
f"unknown action '{action}' for service '{service}'. "
|
||||
f"Known actions: {', '.join(known) or '(none)'}"
|
||||
)
|
||||
|
||||
if method:
|
||||
m = method.lower()
|
||||
picked = [c for c in candidates if c[0].method == m]
|
||||
if not picked:
|
||||
have = sorted({c[0].method.upper() for c in candidates})
|
||||
raise KeyError(
|
||||
f"{service}/{action} does not accept {m.upper()}; "
|
||||
f"available: {', '.join(have)}"
|
||||
)
|
||||
elif len(candidates) == 1:
|
||||
picked = candidates
|
||||
else:
|
||||
prefer = "post" if has_params else "get"
|
||||
picked = [c for c in candidates if c[0].method == prefer] or candidates
|
||||
|
||||
op, bound = picked[0]
|
||||
path = op.path
|
||||
for name, value in bound.items():
|
||||
path = path.replace("{" + name + "}", value)
|
||||
return Route(op.method.upper(), path, bound)
|
||||
|
||||
|
||||
def cache_file(base: str) -> Path:
|
||||
"""Cache location for a base URL. Distinct hosts must not share a file."""
|
||||
slug = base.split("://", 1)[-1].replace("/", "_").replace(":", "_")
|
||||
return Path.home() / ".hanzo" / "cache" / f"openapi-{slug}.json"
|
||||
|
||||
|
||||
def spec_url(base: str) -> str:
|
||||
"""Where to read the registry from.
|
||||
|
||||
Normally the target itself publishes it. ``HANZO_OPENAPI_URL`` separates the
|
||||
two so a catalog can be read from one cloud while calls go to another — the
|
||||
case when driving a partial local host, which mounts only some apps and so
|
||||
cannot serve the whole registry.
|
||||
"""
|
||||
return os.environ.get("HANZO_OPENAPI_URL") or (base.rstrip("/") + SPEC_PATH)
|
||||
|
||||
|
||||
def fetch(base: str, timeout: float = 60.0) -> dict[str, Any]:
|
||||
"""Fetch the spec straight from the running cloud."""
|
||||
url = spec_url(base)
|
||||
req = urllib.request.Request(url, headers={"Accept": "application/json"})
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
return json.loads(resp.read().decode("utf-8"))
|
||||
except (urllib.error.URLError, TimeoutError, ValueError, OSError) as e:
|
||||
raise RuntimeError(f"cannot fetch OpenAPI spec from {url}: {e}") from e
|
||||
|
||||
|
||||
def load(
|
||||
base: str | None = None,
|
||||
refresh: bool = False,
|
||||
ttl: float = CACHE_TTL,
|
||||
) -> tuple[dict[str, Any], str]:
|
||||
"""Return (spec, source). Cache is a speed-up, never a source of truth.
|
||||
|
||||
A stale cache is still preferred over failing the call outright: the
|
||||
catalog only names routes, and a name that has since moved fails loudly at
|
||||
call time anyway.
|
||||
"""
|
||||
base = (base or api_base()).rstrip("/")
|
||||
path = cache_file(spec_url(base))
|
||||
fresh = (
|
||||
not refresh
|
||||
and path.exists()
|
||||
and (time.time() - path.stat().st_mtime) < ttl
|
||||
)
|
||||
if fresh:
|
||||
try:
|
||||
return json.loads(path.read_text()), f"cache:{path}"
|
||||
except (OSError, ValueError):
|
||||
pass # corrupt cache: fall through and refetch
|
||||
|
||||
try:
|
||||
spec = fetch(base)
|
||||
except RuntimeError:
|
||||
if path.exists():
|
||||
try:
|
||||
return json.loads(path.read_text()), f"cache-stale:{path}"
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
raise
|
||||
|
||||
try:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(spec))
|
||||
except OSError:
|
||||
pass # a read-only home must not break the call
|
||||
return spec, spec_url(base)
|
||||
|
||||
|
||||
def split_params(
|
||||
route: Route, params: dict[str, Any]
|
||||
) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
|
||||
"""Split caller params into (query, body) for a route.
|
||||
|
||||
Params already consumed by the path template are dropped so they are not
|
||||
also sent as query/body noise.
|
||||
"""
|
||||
rest: dict[str, Any] = {
|
||||
k: v for k, v in params.items() if k not in route.bound and v is not None
|
||||
}
|
||||
if route.method.lower() in BODY_METHODS:
|
||||
return None, rest
|
||||
return rest or None, None
|
||||
|
||||
|
||||
def iter_actions(ops: Iterable[Operation]) -> list[str]:
|
||||
return sorted({o.action for o in ops})
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "hanzo-tools-api"
|
||||
version = "0.3.1"
|
||||
version = "0.3.2"
|
||||
description = "Generic API tool for calling any REST API via OpenAPI specs - search, explore, and dynamically use ANY API"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
|
||||
@@ -8,6 +8,21 @@ from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# A miniature of what cloud publishes: product tags as services, paths as
|
||||
# actions, one templated path. Enough to pin the whole projection.
|
||||
SPEC_FIXTURE = {
|
||||
"paths": {
|
||||
"/v1/billing/plans": {"get": {"tags": ["billing"]}},
|
||||
"/v1/billing/balance": {"get": {"tags": ["billing"]}},
|
||||
"/v1/kb/search": {"post": {"tags": ["kb"]}},
|
||||
"/v1/iam/{user}": {"get": {"tags": ["iam"]}},
|
||||
"/v1/iam/keys": {
|
||||
"get": {"tags": ["iam"]},
|
||||
"post": {"tags": ["iam"], "requestBody": {}},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
from hanzo_tools.api import (
|
||||
ENV_VAR_MAPPINGS,
|
||||
PROVIDER_CONFIGS,
|
||||
@@ -415,34 +430,85 @@ class TestAPITool:
|
||||
|
||||
|
||||
class TestHanzoTool:
|
||||
"""Tests for unified HanzoTool surface."""
|
||||
"""Tests for the unified HanzoTool surface.
|
||||
|
||||
The surface comes from the OpenAPI registry, so these drive a stub catalog
|
||||
rather than the network: the mapping is what is under test.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _tool():
|
||||
from hanzo_tools.api import spec as openapi
|
||||
|
||||
tool = HanzoTool()
|
||||
tool._catalog = openapi.Catalog(SPEC_FIXTURE)
|
||||
tool._source = "test"
|
||||
return tool
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_services_listing(self):
|
||||
"""Service discovery should return consolidated service list."""
|
||||
tool = HanzoTool()
|
||||
ctx = AsyncMock()
|
||||
result = await tool.call(ctx, service="services")
|
||||
payload = json.loads(result)
|
||||
assert "services" in payload
|
||||
assert "hanzo" not in payload["services"] # service router, not a nested service
|
||||
assert "commerce" in payload["services"]
|
||||
assert "iam" in payload["services"]
|
||||
async def test_services_listing_comes_from_spec(self):
|
||||
tool = self._tool()
|
||||
payload = json.loads(await tool.call(AsyncMock(), service="services"))
|
||||
assert payload["count"] == 3
|
||||
assert set(payload["services"]) == {"billing", "iam", "kb"}
|
||||
assert payload["services"]["billing"]["operations"] == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_args_json(self):
|
||||
"""Invalid JSON args should return structured error."""
|
||||
tool = HanzoTool()
|
||||
ctx = AsyncMock()
|
||||
result = await tool.call(
|
||||
ctx,
|
||||
service="iam",
|
||||
action="users",
|
||||
args="{invalid-json",
|
||||
async def test_empty_action_lists_service_actions(self):
|
||||
tool = self._tool()
|
||||
payload = json.loads(await tool.call(AsyncMock(), service="billing"))
|
||||
assert sorted(payload["actions"]) == ["balance", "plans"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_params_json(self):
|
||||
tool = self._tool()
|
||||
payload = json.loads(
|
||||
await tool.call(AsyncMock(), service="kb", action="search", params="{bad")
|
||||
)
|
||||
payload = json.loads(result)
|
||||
assert "error" in payload and payload["service"] == "kb"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_service_hints_at_listing(self):
|
||||
tool = self._tool()
|
||||
payload = json.loads(await tool.call(AsyncMock(), service="nope", action="x"))
|
||||
assert "error" in payload
|
||||
assert payload["service"] == "iam"
|
||||
assert payload["hint"] == 'call hanzo(service="services") to list services'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_dispatches_to_resolved_route(self):
|
||||
"""A named action must reach exactly the route the spec declares."""
|
||||
tool = self._tool()
|
||||
seen = {}
|
||||
|
||||
class FakeCloud:
|
||||
async def call(self, method, path, params=None, json_body=None):
|
||||
seen.update(method=method, path=path, params=params, body=json_body)
|
||||
return {"ok": True}
|
||||
|
||||
tool._cloud = FakeCloud()
|
||||
payload = json.loads(
|
||||
await tool.call(
|
||||
AsyncMock(), service="kb", action="search", params='{"query":"oauth"}'
|
||||
)
|
||||
)
|
||||
assert seen["method"] == "POST" and seen["path"] == "/v1/kb/search"
|
||||
assert seen["body"] == {"query": "oauth"} # POST carries params as a body
|
||||
assert payload["data"] == {"ok": True}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_path_template_is_filled_from_the_action(self):
|
||||
tool = self._tool()
|
||||
seen = {}
|
||||
|
||||
class FakeCloud:
|
||||
async def call(self, method, path, params=None, json_body=None):
|
||||
seen.update(path=path, params=params)
|
||||
return {}
|
||||
|
||||
tool._cloud = FakeCloud()
|
||||
await tool.call(AsyncMock(), service="iam", action="alice")
|
||||
assert seen["path"] == "/v1/iam/alice"
|
||||
assert not seen["params"] # the bound param is not resent as a query
|
||||
|
||||
|
||||
class TestIntegration:
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
"""Tests for the OpenAPI -> service/action projection.
|
||||
|
||||
Pure mapping tests: no network. These pin the contract that lets one MCP tool
|
||||
stand in for every cloud service, so a regression here silently mis-routes
|
||||
calls rather than failing loudly.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from hanzo_tools.api import spec as openapi
|
||||
|
||||
SPEC = {
|
||||
"paths": {
|
||||
"/v1/billing/plans": {"get": {"tags": ["billing"]}},
|
||||
"/v1/billing/spend-alerts/{id}": {"delete": {"tags": ["billing"]}},
|
||||
"/v1/code/ask": {"get": {"tags": ["code"]}, "post": {"tags": ["code"]}},
|
||||
"/v1/kb/search": {"post": {"tags": ["kb"], "responses": {"200": {}}}},
|
||||
"/v1/iam": {"get": {"tags": ["iam"]}},
|
||||
"/v1/vector/{name}": {"get": {"tags": ["vector"]}},
|
||||
"/v1/vector/stats": {"get": {"tags": ["vector"]}},
|
||||
"/healthz": {"get": {}}, # untagged: still addressable
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def catalog():
|
||||
return openapi.Catalog(SPEC)
|
||||
|
||||
|
||||
class TestActionNaming:
|
||||
@pytest.mark.parametrize(
|
||||
"path,tag,expected",
|
||||
[
|
||||
("/v1/billing/plans", "billing", "plans"),
|
||||
("/v1/iam", "iam", ""),
|
||||
("/v1/vector/{name}", "vector", "{name}"),
|
||||
("/v1/billing/spend-alerts/{id}", "billing", "spend-alerts/{id}"),
|
||||
("/healthz", "_untagged", "healthz"),
|
||||
],
|
||||
)
|
||||
def test_action_of(self, path, tag, expected):
|
||||
assert openapi.action_of(path, tag) == expected
|
||||
|
||||
|
||||
class TestCatalog:
|
||||
def test_services_are_tags(self, catalog):
|
||||
assert catalog.services == [
|
||||
"_untagged",
|
||||
"billing",
|
||||
"code",
|
||||
"iam",
|
||||
"kb",
|
||||
"vector",
|
||||
]
|
||||
|
||||
def test_summary_counts_typed_ops(self, catalog):
|
||||
summary = catalog.summary()
|
||||
assert summary["kb"]["typed"] == 1 # has responses
|
||||
assert summary["billing"]["typed"] == 0 # router-shape only
|
||||
assert summary["code"]["operations"] == 2
|
||||
|
||||
def test_describe_groups_methods_per_action(self, catalog):
|
||||
actions = catalog.describe("code")["actions"]
|
||||
assert actions["ask"]["methods"] == ["GET", "POST"]
|
||||
|
||||
def test_unknown_service_raises(self, catalog):
|
||||
with pytest.raises(KeyError):
|
||||
catalog.resolve("nope", "x")
|
||||
|
||||
def test_unknown_action_lists_known_ones(self, catalog):
|
||||
with pytest.raises(KeyError, match="plans"):
|
||||
catalog.resolve("billing", "nope")
|
||||
|
||||
|
||||
class TestResolution:
|
||||
def test_exact_action(self, catalog):
|
||||
assert catalog.resolve("billing", "plans") == openapi.Route(
|
||||
"GET", "/v1/billing/plans", {}
|
||||
)
|
||||
|
||||
def test_root_action(self, catalog):
|
||||
assert catalog.resolve("iam", "").path == "/v1/iam"
|
||||
|
||||
def test_template_binds_from_action(self, catalog):
|
||||
route = catalog.resolve("vector", "my-coll")
|
||||
assert route.path == "/v1/vector/my-coll"
|
||||
assert route.bound == {"name": "my-coll"}
|
||||
|
||||
def test_literal_beats_template(self, catalog):
|
||||
"""`stats` is a real route; it must not be swallowed by {name}."""
|
||||
route = catalog.resolve("vector", "stats")
|
||||
assert route.path == "/v1/vector/stats" and route.bound == {}
|
||||
|
||||
def test_nested_template(self, catalog):
|
||||
route = catalog.resolve("billing", "spend-alerts/a1")
|
||||
assert route.path == "/v1/billing/spend-alerts/a1"
|
||||
|
||||
def test_method_override(self, catalog):
|
||||
assert catalog.resolve("code", "ask", method="post").method == "POST"
|
||||
|
||||
def test_bad_method_reports_available(self, catalog):
|
||||
with pytest.raises(KeyError, match="GET"):
|
||||
catalog.resolve("billing", "plans", method="delete")
|
||||
|
||||
def test_params_imply_write_absence_implies_read(self, catalog):
|
||||
assert catalog.resolve("code", "ask", has_params=True).method == "POST"
|
||||
assert catalog.resolve("code", "ask", has_params=False).method == "GET"
|
||||
|
||||
|
||||
class TestParamSplit:
|
||||
def test_body_methods_get_a_body(self):
|
||||
route = openapi.Route("POST", "/v1/kb/search", {})
|
||||
query, body = openapi.split_params(route, {"query": "x"})
|
||||
assert query is None and body == {"query": "x"}
|
||||
|
||||
def test_read_methods_get_a_query(self):
|
||||
route = openapi.Route("GET", "/v1/billing/plans", {})
|
||||
query, body = openapi.split_params(route, {"limit": 5})
|
||||
assert query == {"limit": 5} and body is None
|
||||
|
||||
def test_bound_path_params_are_not_resent(self):
|
||||
route = openapi.Route("GET", "/v1/vector/c1", {"name": "c1"})
|
||||
query, _ = openapi.split_params(route, {"name": "c1", "limit": 2})
|
||||
assert query == {"limit": 2}
|
||||
|
||||
|
||||
class TestSpecSource:
|
||||
def test_default_is_the_target_itself(self, monkeypatch):
|
||||
monkeypatch.delenv("HANZO_OPENAPI_URL", raising=False)
|
||||
assert openapi.spec_url("https://api.hanzo.ai") == (
|
||||
"https://api.hanzo.ai/v1/openapi.json"
|
||||
)
|
||||
|
||||
def test_override_decouples_catalog_from_target(self, monkeypatch):
|
||||
monkeypatch.setenv("HANZO_OPENAPI_URL", "https://api.hanzo.ai/v1/openapi.json")
|
||||
assert openapi.spec_url("http://127.0.0.1:18080").startswith("https://")
|
||||
|
||||
def test_cache_file_is_per_source(self, monkeypatch):
|
||||
monkeypatch.delenv("HANZO_OPENAPI_URL", raising=False)
|
||||
a = openapi.cache_file("https://api.hanzo.ai")
|
||||
b = openapi.cache_file("http://127.0.0.1:18080")
|
||||
assert a != b
|
||||
|
||||
def test_stale_cache_beats_failing_the_call(self, monkeypatch, tmp_path):
|
||||
"""A catalog is only route names; a stale one still routes."""
|
||||
cached = tmp_path / "openapi.json"
|
||||
cached.write_text(json.dumps(SPEC))
|
||||
monkeypatch.setattr(openapi, "cache_file", lambda _: cached)
|
||||
monkeypatch.setattr(
|
||||
openapi, "fetch", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("down"))
|
||||
)
|
||||
document, source = openapi.load("https://api.hanzo.ai", refresh=True)
|
||||
assert document == SPEC and source.startswith("cache-stale:")
|
||||
@@ -1,15 +1,12 @@
|
||||
"""Vector/embedding tools for Hanzo AI.
|
||||
|
||||
Default surface is the cloud-backed `VectorTool` (real Qdrant + zen embeddings
|
||||
via api.hanzo.ai). Actions: search, index, embed.
|
||||
The only surface is the cloud-backed `VectorTool`: real zen embeddings into the
|
||||
in-cluster Qdrant via api.hanzo.ai. Actions: search, index, embed.
|
||||
|
||||
The local Infinity store (infinity_store / vector_search / vector_index) remains
|
||||
importable for offline use but is NOT the default and never substitutes a mock
|
||||
silently — the random-vector mock is opt-in via HANZO_VECTOR_ALLOW_MOCK=1.
|
||||
|
||||
Install:
|
||||
pip install hanzo-tools-vector # cloud (default)
|
||||
pip install hanzo-tools-vector[full] # + local Infinity store
|
||||
There is no local store. The former embedded-Infinity path shipped no embedder,
|
||||
so it could only ever produce random vectors — which rank by noise and silently
|
||||
lie to the caller. A tool that lies is worse than a missing tool, so it is gone
|
||||
rather than flag-guarded: every vector here is computed by the real service.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
@@ -1,500 +0,0 @@
|
||||
"""AST analysis and symbol extraction for code understanding."""
|
||||
|
||||
import ast
|
||||
import hashlib
|
||||
from typing import Any, Dict, List, Optional
|
||||
from pathlib import Path
|
||||
from dataclasses import asdict, dataclass
|
||||
|
||||
try:
|
||||
import tree_sitter
|
||||
import tree_sitter_python as tspython
|
||||
|
||||
TREE_SITTER_AVAILABLE = True
|
||||
except ImportError:
|
||||
TREE_SITTER_AVAILABLE = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class Symbol:
|
||||
"""Represents a code symbol (function, class, variable, etc.)."""
|
||||
|
||||
name: str
|
||||
type: str # function, class, variable, import, etc.
|
||||
file_path: str
|
||||
line_start: int
|
||||
line_end: int
|
||||
column_start: int
|
||||
column_end: int
|
||||
scope: str # global, class, function
|
||||
parent: Optional[str] = None # parent class/function
|
||||
docstring: Optional[str] = None
|
||||
signature: Optional[str] = None
|
||||
references: List[str] = None # Files that reference this symbol
|
||||
|
||||
def __post_init__(self):
|
||||
if self.references is None:
|
||||
self.references = []
|
||||
|
||||
|
||||
@dataclass
|
||||
class ASTNode:
|
||||
"""Represents an AST node with metadata."""
|
||||
|
||||
type: str
|
||||
name: Optional[str]
|
||||
line_start: int
|
||||
line_end: int
|
||||
column_start: int
|
||||
column_end: int
|
||||
children: List["ASTNode"] = None
|
||||
parent: Optional[str] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.children is None:
|
||||
self.children = []
|
||||
|
||||
|
||||
@dataclass
|
||||
class FileAST:
|
||||
"""Complete AST representation of a file."""
|
||||
|
||||
file_path: str
|
||||
file_hash: str
|
||||
language: str
|
||||
symbols: List[Symbol]
|
||||
ast_nodes: List[ASTNode]
|
||||
imports: List[str]
|
||||
exports: List[str]
|
||||
dependencies: List[str] # Files this file depends on
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for storage."""
|
||||
return {
|
||||
"file_path": self.file_path,
|
||||
"file_hash": self.file_hash,
|
||||
"language": self.language,
|
||||
"symbols": [asdict(s) for s in self.symbols],
|
||||
"ast_nodes": [asdict(n) for n in self.ast_nodes],
|
||||
"imports": self.imports,
|
||||
"exports": self.exports,
|
||||
"dependencies": self.dependencies,
|
||||
}
|
||||
|
||||
|
||||
class ASTAnalyzer:
|
||||
"""Analyzes code files and extracts AST information and symbols."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the AST analyzer."""
|
||||
self.parsers = {}
|
||||
self._setup_parsers()
|
||||
|
||||
def _setup_parsers(self):
|
||||
"""Set up tree-sitter parsers for different languages."""
|
||||
if TREE_SITTER_AVAILABLE:
|
||||
try:
|
||||
# Python parser
|
||||
self.parsers["python"] = tree_sitter.Language(tspython.language())
|
||||
except Exception as e:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning(f"Could not initialize Python parser: {e}")
|
||||
|
||||
def analyze_file(self, file_path: str) -> Optional[FileAST]:
|
||||
"""Analyze a file and extract AST information and symbols.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file to analyze
|
||||
|
||||
Returns:
|
||||
FileAST object with extracted information, or None if analysis fails
|
||||
"""
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
return None
|
||||
|
||||
# Determine language
|
||||
language = self._detect_language(path)
|
||||
if not language:
|
||||
return None
|
||||
|
||||
try:
|
||||
# Read file content
|
||||
content = path.read_text(encoding="utf-8")
|
||||
file_hash = hashlib.sha256(content.encode()).hexdigest()
|
||||
|
||||
# Extract symbols and AST
|
||||
if language == "python":
|
||||
return self._analyze_python_file(file_path, content, file_hash)
|
||||
else:
|
||||
# Generic analysis for other languages
|
||||
return self._analyze_generic_file(
|
||||
file_path, content, file_hash, language
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.error(f"Error analyzing file {file_path}: {e}")
|
||||
return None
|
||||
|
||||
def _detect_language(self, path: Path) -> Optional[str]:
|
||||
"""Detect programming language from file extension."""
|
||||
extension = path.suffix.lower()
|
||||
|
||||
language_map = {
|
||||
".py": "python",
|
||||
".js": "javascript",
|
||||
".ts": "typescript",
|
||||
".jsx": "javascript",
|
||||
".tsx": "typescript",
|
||||
".java": "java",
|
||||
".cpp": "cpp",
|
||||
".c": "c",
|
||||
".h": "c",
|
||||
".hpp": "cpp",
|
||||
".rs": "rust",
|
||||
".go": "go",
|
||||
".rb": "ruby",
|
||||
".php": "php",
|
||||
".cs": "csharp",
|
||||
".swift": "swift",
|
||||
".kt": "kotlin",
|
||||
".scala": "scala",
|
||||
".clj": "clojure",
|
||||
".hs": "haskell",
|
||||
".ml": "ocaml",
|
||||
".elm": "elm",
|
||||
".dart": "dart",
|
||||
".lua": "lua",
|
||||
".r": "r",
|
||||
".m": "objective-c",
|
||||
".mm": "objective-cpp",
|
||||
}
|
||||
|
||||
return language_map.get(extension)
|
||||
|
||||
def _analyze_python_file(
|
||||
self, file_path: str, content: str, file_hash: str
|
||||
) -> FileAST:
|
||||
"""Analyze Python file using both AST and tree-sitter."""
|
||||
symbols = []
|
||||
ast_nodes = []
|
||||
imports = []
|
||||
exports = []
|
||||
dependencies = []
|
||||
|
||||
try:
|
||||
# Parse with Python AST
|
||||
tree = ast.parse(content)
|
||||
|
||||
# Extract symbols using AST visitor
|
||||
visitor = PythonSymbolExtractor(file_path)
|
||||
visitor.visit(tree)
|
||||
|
||||
symbols.extend(visitor.symbols)
|
||||
imports.extend(visitor.imports)
|
||||
exports.extend(visitor.exports)
|
||||
dependencies.extend(visitor.dependencies)
|
||||
|
||||
# If tree-sitter is available, get more detailed AST
|
||||
if TREE_SITTER_AVAILABLE and "python" in self.parsers:
|
||||
parser = tree_sitter.Parser(self.parsers["python"])
|
||||
ts_tree = parser.parse(content.encode())
|
||||
ast_nodes = self._extract_tree_sitter_nodes(ts_tree.root_node, content)
|
||||
|
||||
except SyntaxError as e:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.error(f"Syntax error in {file_path}: {e}")
|
||||
except Exception as e:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.error(f"Error parsing Python file {file_path}: {e}")
|
||||
|
||||
return FileAST(
|
||||
file_path=file_path,
|
||||
file_hash=file_hash,
|
||||
language="python",
|
||||
symbols=symbols,
|
||||
ast_nodes=ast_nodes,
|
||||
imports=imports,
|
||||
exports=exports,
|
||||
dependencies=dependencies,
|
||||
)
|
||||
|
||||
def _analyze_generic_file(
|
||||
self, file_path: str, content: str, file_hash: str, language: str
|
||||
) -> FileAST:
|
||||
"""Generic analysis for non-Python files."""
|
||||
# For now, just basic line-based analysis
|
||||
# Could be enhanced with language-specific parsers
|
||||
|
||||
symbols = []
|
||||
ast_nodes = []
|
||||
imports = []
|
||||
exports = []
|
||||
dependencies = []
|
||||
|
||||
# Basic pattern matching for common constructs
|
||||
lines = content.split("\n")
|
||||
for i, line in enumerate(lines, 1):
|
||||
line = line.strip()
|
||||
|
||||
# Basic function detection (works for many C-style languages)
|
||||
if language in ["javascript", "typescript", "java", "cpp", "c"]:
|
||||
if (
|
||||
"function " in line
|
||||
or line.startswith("def ")
|
||||
or " function(" in line
|
||||
):
|
||||
# Extract function name
|
||||
parts = line.split()
|
||||
for j, part in enumerate(parts):
|
||||
if part == "function" and j + 1 < len(parts):
|
||||
func_name = parts[j + 1].split("(")[0]
|
||||
symbols.append(
|
||||
Symbol(
|
||||
name=func_name,
|
||||
type="function",
|
||||
file_path=file_path,
|
||||
line_start=i,
|
||||
line_end=i,
|
||||
column_start=0,
|
||||
column_end=len(line),
|
||||
scope="global",
|
||||
)
|
||||
)
|
||||
break
|
||||
|
||||
# Basic import detection
|
||||
if "import " in line or "#include " in line or "require(" in line:
|
||||
imports.append(line)
|
||||
|
||||
return FileAST(
|
||||
file_path=file_path,
|
||||
file_hash=file_hash,
|
||||
language=language,
|
||||
symbols=symbols,
|
||||
ast_nodes=ast_nodes,
|
||||
imports=imports,
|
||||
exports=exports,
|
||||
dependencies=dependencies,
|
||||
)
|
||||
|
||||
def _extract_tree_sitter_nodes(self, node, content: str) -> List[ASTNode]:
|
||||
"""Extract AST nodes from tree-sitter parse tree."""
|
||||
nodes = []
|
||||
|
||||
def traverse(ts_node, parent_name=None):
|
||||
node_name = None
|
||||
|
||||
# Try to extract node name for named nodes
|
||||
if ts_node.type in [
|
||||
"function_definition",
|
||||
"class_definition",
|
||||
"identifier",
|
||||
]:
|
||||
for child in ts_node.children:
|
||||
if child.type == "identifier":
|
||||
start_byte = child.start_byte
|
||||
end_byte = child.end_byte
|
||||
node_name = content[start_byte:end_byte]
|
||||
break
|
||||
|
||||
ast_node = ASTNode(
|
||||
type=ts_node.type,
|
||||
name=node_name,
|
||||
line_start=ts_node.start_point[0] + 1,
|
||||
line_end=ts_node.end_point[0] + 1,
|
||||
column_start=ts_node.start_point[1],
|
||||
column_end=ts_node.end_point[1],
|
||||
parent=parent_name,
|
||||
)
|
||||
|
||||
nodes.append(ast_node)
|
||||
|
||||
# Recursively process children
|
||||
for child in ts_node.children:
|
||||
traverse(child, node_name or parent_name)
|
||||
|
||||
traverse(node)
|
||||
return nodes
|
||||
|
||||
|
||||
class PythonSymbolExtractor(ast.NodeVisitor):
|
||||
"""AST visitor for extracting Python symbols."""
|
||||
|
||||
def __init__(self, file_path: str):
|
||||
self.file_path = file_path
|
||||
self.symbols = []
|
||||
self.imports = []
|
||||
self.exports = []
|
||||
self.dependencies = []
|
||||
self.scope_stack = ["global"]
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
"""Visit function definitions."""
|
||||
scope = ".".join(self.scope_stack)
|
||||
parent = self.scope_stack[-1] if len(self.scope_stack) > 1 else None
|
||||
|
||||
# Extract docstring
|
||||
docstring = None
|
||||
if (
|
||||
node.body
|
||||
and isinstance(node.body[0], ast.Expr)
|
||||
and isinstance(node.body[0].value, ast.Constant)
|
||||
and isinstance(node.body[0].value.value, str)
|
||||
):
|
||||
docstring = node.body[0].value.value
|
||||
|
||||
# Create function signature
|
||||
args = [arg.arg for arg in node.args.args]
|
||||
signature = f"{node.name}({', '.join(args)})"
|
||||
|
||||
symbol = Symbol(
|
||||
name=node.name,
|
||||
type="function",
|
||||
file_path=self.file_path,
|
||||
line_start=node.lineno,
|
||||
line_end=node.end_lineno or node.lineno,
|
||||
column_start=node.col_offset,
|
||||
column_end=node.end_col_offset or 0,
|
||||
scope=scope,
|
||||
parent=parent if parent != "global" else None,
|
||||
docstring=docstring,
|
||||
signature=signature,
|
||||
)
|
||||
|
||||
self.symbols.append(symbol)
|
||||
|
||||
# Enter function scope
|
||||
self.scope_stack.append(node.name)
|
||||
self.generic_visit(node)
|
||||
self.scope_stack.pop()
|
||||
|
||||
def visit_AsyncFunctionDef(self, node):
|
||||
"""Visit async function definitions."""
|
||||
self.visit_FunctionDef(node) # Same logic
|
||||
|
||||
def visit_ClassDef(self, node):
|
||||
"""Visit class definitions."""
|
||||
scope = ".".join(self.scope_stack)
|
||||
parent = self.scope_stack[-1] if len(self.scope_stack) > 1 else None
|
||||
|
||||
# Extract docstring
|
||||
docstring = None
|
||||
if (
|
||||
node.body
|
||||
and isinstance(node.body[0], ast.Expr)
|
||||
and isinstance(node.body[0].value, ast.Constant)
|
||||
and isinstance(node.body[0].value.value, str)
|
||||
):
|
||||
docstring = node.body[0].value.value
|
||||
|
||||
# Extract base classes
|
||||
bases = [self._get_name(base) for base in node.bases]
|
||||
signature = (
|
||||
f"class {node.name}({', '.join(bases)})" if bases else f"class {node.name}"
|
||||
)
|
||||
|
||||
symbol = Symbol(
|
||||
name=node.name,
|
||||
type="class",
|
||||
file_path=self.file_path,
|
||||
line_start=node.lineno,
|
||||
line_end=node.end_lineno or node.lineno,
|
||||
column_start=node.col_offset,
|
||||
column_end=node.end_col_offset or 0,
|
||||
scope=scope,
|
||||
parent=parent if parent != "global" else None,
|
||||
docstring=docstring,
|
||||
signature=signature,
|
||||
)
|
||||
|
||||
self.symbols.append(symbol)
|
||||
|
||||
# Enter class scope
|
||||
self.scope_stack.append(node.name)
|
||||
self.generic_visit(node)
|
||||
self.scope_stack.pop()
|
||||
|
||||
def visit_Import(self, node):
|
||||
"""Visit import statements."""
|
||||
for alias in node.names:
|
||||
import_name = alias.name
|
||||
self.imports.append(import_name)
|
||||
if "." not in import_name: # Top-level module
|
||||
self.dependencies.append(import_name)
|
||||
|
||||
def visit_ImportFrom(self, node):
|
||||
"""Visit from...import statements."""
|
||||
if node.module:
|
||||
self.imports.append(node.module)
|
||||
if "." not in node.module: # Top-level module
|
||||
self.dependencies.append(node.module)
|
||||
|
||||
for alias in node.names:
|
||||
if alias.name != "*":
|
||||
import_item = (
|
||||
f"{node.module}.{alias.name}" if node.module else alias.name
|
||||
)
|
||||
self.imports.append(import_item)
|
||||
|
||||
def visit_Assign(self, node):
|
||||
"""Visit variable assignments."""
|
||||
scope = ".".join(self.scope_stack)
|
||||
parent = self.scope_stack[-1] if len(self.scope_stack) > 1 else None
|
||||
|
||||
for target in node.targets:
|
||||
if isinstance(target, ast.Name):
|
||||
symbol = Symbol(
|
||||
name=target.id,
|
||||
type="variable",
|
||||
file_path=self.file_path,
|
||||
line_start=node.lineno,
|
||||
line_end=node.end_lineno or node.lineno,
|
||||
column_start=node.col_offset,
|
||||
column_end=node.end_col_offset or 0,
|
||||
scope=scope,
|
||||
parent=parent if parent != "global" else None,
|
||||
)
|
||||
self.symbols.append(symbol)
|
||||
|
||||
self.generic_visit(node)
|
||||
|
||||
def _get_name(self, node):
|
||||
"""Extract name from AST node."""
|
||||
if isinstance(node, ast.Name):
|
||||
return node.id
|
||||
elif isinstance(node, ast.Attribute):
|
||||
return f"{self._get_name(node.value)}.{node.attr}"
|
||||
elif isinstance(node, ast.Constant):
|
||||
return str(node.value)
|
||||
else:
|
||||
return str(node)
|
||||
|
||||
|
||||
def create_symbol_embedding_text(symbol: Symbol) -> str:
|
||||
"""Create text representation of symbol for vector embedding."""
|
||||
parts = [
|
||||
f"Symbol: {symbol.name}",
|
||||
f"Type: {symbol.type}",
|
||||
f"Scope: {symbol.scope}",
|
||||
]
|
||||
|
||||
if symbol.parent:
|
||||
parts.append(f"Parent: {symbol.parent}")
|
||||
|
||||
if symbol.signature:
|
||||
parts.append(f"Signature: {symbol.signature}")
|
||||
|
||||
if symbol.docstring:
|
||||
parts.append(f"Documentation: {symbol.docstring}")
|
||||
|
||||
return " | ".join(parts)
|
||||
@@ -9,9 +9,9 @@ Backend map (all verified on api.hanzo.ai /v1):
|
||||
- index → POST /v1/code/index ({repo, files:[{path,content}]})
|
||||
- embed → POST /v1/embeddings ({model, input} → 1024-dim vectors)
|
||||
|
||||
A `collection` names the index namespace (the backend's `repo`). The local
|
||||
Infinity store remains available behind an explicit opt-in (see infinity_store);
|
||||
this tool never falls back to a mock silently.
|
||||
A `collection` names the index namespace (the backend's `repo`). There is no
|
||||
local store and no fallback: if the service is unreachable this fails loudly
|
||||
rather than returning vectors it made up.
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -66,7 +66,7 @@ class VectorTool(BaseTool):
|
||||
"""Unified cloud vector tool: semantic search, index, embed (HIP-0300)."""
|
||||
|
||||
name: ClassVar[str] = "vector"
|
||||
VERSION: ClassVar[str] = "0.2.0"
|
||||
VERSION: ClassVar[str] = "0.2.2"
|
||||
|
||||
def __init__(self, cwd: str | None = None):
|
||||
super().__init__()
|
||||
|
||||
@@ -1,476 +0,0 @@
|
||||
"""Git repository ingester for comprehensive code indexing.
|
||||
|
||||
This module provides functionality to ingest entire git repositories including:
|
||||
- Full git history and commit metadata
|
||||
- File contents at different points in time
|
||||
- AST analysis via tree-sitter
|
||||
- Symbol extraction and cross-references
|
||||
- Blame information for line-level attribution
|
||||
"""
|
||||
|
||||
import logging
|
||||
import subprocess
|
||||
from typing import Any, Dict, List, Optional
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .ast_analyzer import ASTAnalyzer
|
||||
from .infinity_store import InfinityVectorStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GitCommit:
|
||||
"""Represents a git commit."""
|
||||
|
||||
hash: str
|
||||
author: str
|
||||
author_email: str
|
||||
timestamp: int
|
||||
message: str
|
||||
files: List[Dict[str, str]] # [{'status': 'M', 'filename': 'main.py'}]
|
||||
parent_hashes: List[str]
|
||||
|
||||
|
||||
@dataclass
|
||||
class GitFileHistory:
|
||||
"""History of a single file."""
|
||||
|
||||
file_path: str
|
||||
commits: List[GitCommit]
|
||||
current_content: Optional[str]
|
||||
line_blame: Dict[int, Dict[str, Any]] # line_number -> blame info
|
||||
|
||||
|
||||
class GitIngester:
|
||||
"""Ingests git repositories into vector store."""
|
||||
|
||||
def __init__(self, vector_store: InfinityVectorStore):
|
||||
"""Initialize the git ingester.
|
||||
|
||||
Args:
|
||||
vector_store: The vector store to ingest into
|
||||
"""
|
||||
self.vector_store = vector_store
|
||||
self.ast_analyzer = ASTAnalyzer()
|
||||
self._commit_cache: Dict[str, GitCommit] = {}
|
||||
|
||||
def ingest_repository(
|
||||
self,
|
||||
repo_path: str,
|
||||
branch: str = "HEAD",
|
||||
include_history: bool = True,
|
||||
include_diffs: bool = True,
|
||||
include_blame: bool = True,
|
||||
file_patterns: Optional[List[str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Ingest an entire git repository.
|
||||
|
||||
Args:
|
||||
repo_path: Path to the git repository
|
||||
branch: Branch to ingest (default: HEAD)
|
||||
include_history: Whether to include commit history
|
||||
include_diffs: Whether to include diff information
|
||||
include_blame: Whether to include blame information
|
||||
file_patterns: List of file patterns to include (e.g., ["*.py", "*.js"])
|
||||
|
||||
Returns:
|
||||
Summary of ingestion results
|
||||
"""
|
||||
repo_path = Path(repo_path)
|
||||
if not (repo_path / ".git").exists():
|
||||
raise ValueError(f"Not a git repository: {repo_path}")
|
||||
|
||||
logger.info(f"Starting ingestion of repository: {repo_path}")
|
||||
|
||||
results = {
|
||||
"repository": str(repo_path),
|
||||
"branch": branch,
|
||||
"commits_processed": 0,
|
||||
"commits_indexed": 0,
|
||||
"files_indexed": 0,
|
||||
"symbols_extracted": 0,
|
||||
"diffs_indexed": 0,
|
||||
"blame_entries": 0,
|
||||
"errors": [],
|
||||
}
|
||||
|
||||
try:
|
||||
# Get current branch/commit
|
||||
current_commit = self._get_current_commit(repo_path)
|
||||
results["current_commit"] = current_commit
|
||||
|
||||
# Get list of files to process
|
||||
files = self._get_repository_files(repo_path, file_patterns)
|
||||
logger.info(f"Found {len(files)} files to process")
|
||||
|
||||
# Process each file
|
||||
for file_path in files:
|
||||
try:
|
||||
self._process_file(
|
||||
repo_path,
|
||||
file_path,
|
||||
include_history=include_history,
|
||||
include_blame=include_blame,
|
||||
results=results,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing {file_path}: {e}")
|
||||
results["errors"].append(f"{file_path}: {str(e)}")
|
||||
|
||||
# Process commit history if requested
|
||||
if include_history:
|
||||
commits = self._get_commit_history(repo_path, branch)
|
||||
results["commits_processed"] = len(commits)
|
||||
|
||||
for commit in commits:
|
||||
self._index_commit(commit, include_diffs=include_diffs)
|
||||
results["commits_indexed"] = results.get("commits_indexed", 0) + 1
|
||||
|
||||
if include_diffs:
|
||||
results["diffs_indexed"] += len(commit.files)
|
||||
|
||||
# Create repository metadata document
|
||||
self._index_repository_metadata(repo_path, results)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Repository ingestion failed: {e}")
|
||||
results["errors"].append(f"Fatal error: {str(e)}")
|
||||
|
||||
logger.info(f"Ingestion complete: {results}")
|
||||
return results
|
||||
|
||||
def _get_current_commit(self, repo_path: Path) -> str:
|
||||
"""Get the current commit hash."""
|
||||
result = subprocess.run(
|
||||
["git", "rev-parse", "HEAD"],
|
||||
cwd=repo_path,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
return result.stdout.strip()
|
||||
|
||||
def _get_repository_files(
|
||||
self, repo_path: Path, patterns: Optional[List[str]] = None
|
||||
) -> List[Path]:
|
||||
"""Get list of files in repository matching patterns."""
|
||||
# Use git ls-files to respect .gitignore
|
||||
cmd = ["git", "ls-files"]
|
||||
|
||||
result = subprocess.run(
|
||||
cmd, cwd=repo_path, capture_output=True, text=True, check=True
|
||||
)
|
||||
|
||||
files = []
|
||||
for line in result.stdout.strip().split("\n"):
|
||||
if line:
|
||||
file_path = repo_path / line
|
||||
if file_path.exists():
|
||||
# Apply pattern filtering if specified
|
||||
if patterns:
|
||||
if any(file_path.match(pattern) for pattern in patterns):
|
||||
files.append(file_path)
|
||||
else:
|
||||
files.append(file_path)
|
||||
|
||||
return files
|
||||
|
||||
def _get_commit_history(
|
||||
self, repo_path: Path, branch: str = "HEAD", max_commits: int = 1000
|
||||
) -> List[GitCommit]:
|
||||
"""Get commit history for the repository."""
|
||||
# Get commit list with basic info
|
||||
result = subprocess.run(
|
||||
[
|
||||
"git",
|
||||
"log",
|
||||
branch,
|
||||
f"--max-count={max_commits}",
|
||||
"--pretty=format:%H|%P|%an|%ae|%at|%s",
|
||||
],
|
||||
cwd=repo_path,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
|
||||
commits = []
|
||||
for line in result.stdout.strip().split("\n"):
|
||||
if line:
|
||||
parts = line.split("|", 5)
|
||||
if len(parts) >= 6:
|
||||
commit_hash = parts[0]
|
||||
parent_hashes = parts[1].split() if parts[1] else []
|
||||
|
||||
# Get file changes for this commit
|
||||
files = self._get_commit_files(repo_path, commit_hash)
|
||||
|
||||
commit = GitCommit(
|
||||
hash=commit_hash,
|
||||
parent_hashes=parent_hashes,
|
||||
author=parts[2],
|
||||
author_email=parts[3],
|
||||
timestamp=int(parts[4]),
|
||||
message=parts[5],
|
||||
files=files,
|
||||
)
|
||||
commits.append(commit)
|
||||
self._commit_cache[commit_hash] = commit
|
||||
|
||||
return commits
|
||||
|
||||
def _get_commit_files(
|
||||
self, repo_path: Path, commit_hash: str
|
||||
) -> List[Dict[str, str]]:
|
||||
"""Get list of files changed in a commit."""
|
||||
result = subprocess.run(
|
||||
["git", "show", "--name-status", "--format=", commit_hash],
|
||||
cwd=repo_path,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
|
||||
files = []
|
||||
for line in result.stdout.strip().split("\n"):
|
||||
if line and "\t" in line:
|
||||
parts = line.split("\t", 1)
|
||||
if len(parts) == 2:
|
||||
files.append({"status": parts[0], "filename": parts[1]})
|
||||
|
||||
return files
|
||||
|
||||
def _process_file(
|
||||
self,
|
||||
repo_path: Path,
|
||||
file_path: Path,
|
||||
include_history: bool,
|
||||
include_blame: bool,
|
||||
results: Dict[str, Any],
|
||||
):
|
||||
"""Process a single file."""
|
||||
relative_path = file_path.relative_to(repo_path)
|
||||
|
||||
# Read current content
|
||||
try:
|
||||
content = file_path.read_text(encoding="utf-8")
|
||||
except UnicodeDecodeError:
|
||||
content = file_path.read_text(encoding="latin-1")
|
||||
|
||||
# Get file metadata
|
||||
metadata = {
|
||||
"repository": str(repo_path),
|
||||
"relative_path": str(relative_path),
|
||||
"file_type": file_path.suffix,
|
||||
"size": file_path.stat().st_size,
|
||||
}
|
||||
|
||||
# Add git history metadata if requested
|
||||
if include_history:
|
||||
history = self._get_file_history(repo_path, relative_path)
|
||||
metadata["commit_count"] = len(history)
|
||||
if history:
|
||||
metadata["first_commit"] = history[-1]["hash"]
|
||||
metadata["last_commit"] = history[0]["hash"]
|
||||
metadata["last_modified"] = datetime.fromtimestamp(
|
||||
history[0]["timestamp"]
|
||||
).isoformat()
|
||||
|
||||
# Add blame information if requested
|
||||
if include_blame:
|
||||
blame_data = self._get_file_blame(repo_path, relative_path)
|
||||
metadata["unique_authors"] = len(
|
||||
set(b["author"] for b in blame_data.values())
|
||||
)
|
||||
|
||||
# Index the file content
|
||||
doc_ids = self.vector_store.add_file(
|
||||
str(file_path), chunk_size=1000, chunk_overlap=200, metadata=metadata
|
||||
)
|
||||
results["files_indexed"] += 1
|
||||
|
||||
# Perform AST analysis for supported languages
|
||||
if file_path.suffix in [".py", ".js", ".ts", ".java", ".cpp", ".c"]:
|
||||
try:
|
||||
file_ast = self.ast_analyzer.analyze_file(str(file_path))
|
||||
if file_ast:
|
||||
# Store complete AST
|
||||
self.vector_store._store_file_ast(file_ast)
|
||||
|
||||
# Store individual symbols
|
||||
self.vector_store._store_symbols(file_ast.symbols)
|
||||
results["symbols_extracted"] += len(file_ast.symbols)
|
||||
|
||||
# Store cross-references
|
||||
self.vector_store._store_references(file_ast)
|
||||
except Exception as e:
|
||||
logger.warning(f"AST analysis failed for {file_path}: {e}")
|
||||
|
||||
def _get_file_history(
|
||||
self, repo_path: Path, file_path: Path
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Get commit history for a specific file."""
|
||||
result = subprocess.run(
|
||||
[
|
||||
"git",
|
||||
"log",
|
||||
"--follow",
|
||||
"--pretty=format:%H|%at|%an|%s",
|
||||
"--",
|
||||
str(file_path),
|
||||
],
|
||||
cwd=repo_path,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
return []
|
||||
|
||||
history = []
|
||||
for line in result.stdout.strip().split("\n"):
|
||||
if line:
|
||||
parts = line.split("|", 3)
|
||||
if len(parts) >= 4:
|
||||
history.append(
|
||||
{
|
||||
"hash": parts[0],
|
||||
"timestamp": int(parts[1]),
|
||||
"author": parts[2],
|
||||
"message": parts[3],
|
||||
}
|
||||
)
|
||||
|
||||
return history
|
||||
|
||||
def _get_file_blame(
|
||||
self, repo_path: Path, file_path: Path
|
||||
) -> Dict[int, Dict[str, Any]]:
|
||||
"""Get blame information for a file."""
|
||||
result = subprocess.run(
|
||||
["git", "blame", "--line-porcelain", "--", str(file_path)],
|
||||
cwd=repo_path,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
return {}
|
||||
|
||||
blame_data = {}
|
||||
current_commit = None
|
||||
current_line = None
|
||||
author = None
|
||||
timestamp = None
|
||||
|
||||
for line in result.stdout.strip().split("\n"):
|
||||
if line and not line.startswith("\t"):
|
||||
parts = line.split(" ")
|
||||
if len(parts) >= 3 and len(parts[0]) == 40: # SHA-1 hash
|
||||
current_commit = parts[0]
|
||||
current_line = int(parts[2])
|
||||
elif line.startswith("author "):
|
||||
author = line[7:]
|
||||
elif line.startswith("author-time "):
|
||||
timestamp = int(line[12:])
|
||||
|
||||
# We have all the data for this line
|
||||
if current_line and author:
|
||||
blame_data[current_line] = {
|
||||
"commit": current_commit,
|
||||
"author": author,
|
||||
"timestamp": timestamp,
|
||||
}
|
||||
|
||||
return blame_data
|
||||
|
||||
def _index_commit(self, commit: GitCommit, include_diffs: bool = True):
|
||||
"""Index a single commit."""
|
||||
# Create commit document
|
||||
commit_doc = f"""Git Commit: {commit.hash}
|
||||
Author: {commit.author} <{commit.author_email}>
|
||||
Date: {datetime.fromtimestamp(commit.timestamp).isoformat()}
|
||||
Message: {commit.message}
|
||||
|
||||
Files changed: {len(commit.files)}
|
||||
"""
|
||||
|
||||
for file_info in commit.files:
|
||||
commit_doc += f"\n{file_info['status']}\t{file_info['filename']}"
|
||||
|
||||
# Index commit
|
||||
metadata = {
|
||||
"type": "git_commit",
|
||||
"commit_hash": commit.hash,
|
||||
"author": commit.author,
|
||||
"timestamp": commit.timestamp,
|
||||
"file_count": len(commit.files),
|
||||
}
|
||||
|
||||
self.vector_store.add_document(commit_doc, metadata)
|
||||
|
||||
# Index diffs if requested
|
||||
if include_diffs:
|
||||
for file_info in commit.files:
|
||||
self._index_commit_diff(commit, file_info["filename"])
|
||||
|
||||
def _index_commit_diff(self, commit: GitCommit, filename: str):
|
||||
"""Index the diff for a specific file in a commit."""
|
||||
# This is a simplified version - in practice you'd want to
|
||||
# parse the actual diff and store meaningful chunks
|
||||
metadata = {
|
||||
"type": "git_diff",
|
||||
"commit_hash": commit.hash,
|
||||
"filename": filename,
|
||||
"author": commit.author,
|
||||
"timestamp": commit.timestamp,
|
||||
}
|
||||
|
||||
# Create a document representing this change
|
||||
diff_doc = f"""File: {filename}
|
||||
Commit: {commit.hash}
|
||||
Author: {commit.author}
|
||||
Message: {commit.message}
|
||||
"""
|
||||
|
||||
self.vector_store.add_document(diff_doc, metadata)
|
||||
|
||||
def _index_repository_metadata(self, repo_path: Path, results: Dict[str, Any]):
|
||||
"""Index overall repository metadata."""
|
||||
# Get repository info
|
||||
remote_result = subprocess.run(
|
||||
["git", "remote", "get-url", "origin"],
|
||||
cwd=repo_path,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
remote_url = (
|
||||
remote_result.stdout.strip() if remote_result.returncode == 0 else None
|
||||
)
|
||||
|
||||
# Create repository summary document
|
||||
repo_doc = f"""Repository: {repo_path.name}
|
||||
Path: {repo_path}
|
||||
Remote: {remote_url or "No remote"}
|
||||
Current Commit: {results.get("current_commit", "Unknown")}
|
||||
|
||||
Statistics:
|
||||
- Files indexed: {results["files_indexed"]}
|
||||
- Commits processed: {results["commits_processed"]}
|
||||
- Symbols extracted: {results["symbols_extracted"]}
|
||||
- Diffs indexed: {results["diffs_indexed"]}
|
||||
"""
|
||||
|
||||
metadata = {
|
||||
"type": "repository",
|
||||
"name": repo_path.name,
|
||||
"path": str(repo_path),
|
||||
"remote_url": remote_url,
|
||||
**results,
|
||||
}
|
||||
|
||||
self.vector_store.add_document(repo_doc, metadata)
|
||||
@@ -1,416 +0,0 @@
|
||||
"""Index tool for managing vector store indexing."""
|
||||
|
||||
import os
|
||||
import time
|
||||
from typing import Unpack, Annotated, TypedDict, final, override
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import (
|
||||
BaseTool,
|
||||
ToolContext,
|
||||
PermissionManager,
|
||||
auto_timeout,
|
||||
create_tool_context,
|
||||
)
|
||||
|
||||
from .git_ingester import GitIngester
|
||||
from .project_manager import ProjectVectorManager
|
||||
|
||||
Path_str = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Path to index (defaults to current working directory)",
|
||||
min_length=1,
|
||||
),
|
||||
]
|
||||
|
||||
IncludeGitHistory = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Include git history in the index",
|
||||
default=True,
|
||||
),
|
||||
]
|
||||
|
||||
FilePatterns = Annotated[
|
||||
list[str] | None,
|
||||
Field(
|
||||
description="File patterns to include (e.g., ['*.py', '*.js'])",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
ShowStats = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Show detailed statistics after indexing",
|
||||
default=True,
|
||||
),
|
||||
]
|
||||
|
||||
Force = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Force re-indexing even if already indexed",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class IndexToolParams(TypedDict, total=False):
|
||||
"""Parameters for the index tool."""
|
||||
|
||||
path: str
|
||||
include_git_history: bool
|
||||
file_patterns: list[str] | None
|
||||
show_stats: bool
|
||||
force: bool
|
||||
|
||||
|
||||
@final
|
||||
class IndexTool(BaseTool):
|
||||
"""Tool for indexing files and git history into vector store."""
|
||||
|
||||
def __init__(self, permission_manager: PermissionManager):
|
||||
"""Initialize the index tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.project_manager = ProjectVectorManager(permission_manager)
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "index"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Index files and git history into the vector store for semantic search.
|
||||
|
||||
This tool:
|
||||
- Indexes all project files into a vector database
|
||||
- Includes git history (commits, diffs, blame) when available
|
||||
- Supports incremental updates
|
||||
- Shows statistics about indexed content
|
||||
- Automatically creates project-specific databases
|
||||
|
||||
Usage:
|
||||
- index: Index the current directory
|
||||
- index --path /path/to/project: Index a specific path
|
||||
- index --file-patterns "*.py" "*.js": Index only specific file types
|
||||
- index --no-git-history: Skip git history indexing
|
||||
- index --force: Force re-indexing of all files"""
|
||||
|
||||
@override
|
||||
@auto_timeout("index")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[IndexToolParams],
|
||||
) -> str:
|
||||
"""Execute the index tool.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Indexing result and statistics
|
||||
"""
|
||||
start_time = time.time()
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
# Extract parameters
|
||||
path = params.get("path", os.getcwd())
|
||||
include_git_history = params.get("include_git_history", True)
|
||||
file_patterns = params.get("file_patterns")
|
||||
show_stats = params.get("show_stats", True)
|
||||
force = params.get("force", False)
|
||||
|
||||
# Resolve absolute path
|
||||
abs_path = os.path.abspath(path)
|
||||
|
||||
# Check permissions
|
||||
if not self.permission_manager.is_path_allowed(abs_path):
|
||||
return f"Permission denied: {abs_path}"
|
||||
|
||||
# Check if path exists
|
||||
if not os.path.exists(abs_path):
|
||||
return f"Path does not exist: {abs_path}"
|
||||
|
||||
await tool_ctx.info(f"Starting indexing of {abs_path}")
|
||||
|
||||
try:
|
||||
# Get or create vector store for this project
|
||||
vector_store = self.project_manager.get_project_store(abs_path)
|
||||
|
||||
# Check if already indexed (unless force)
|
||||
if not force:
|
||||
stats = await vector_store.get_stats()
|
||||
if stats and stats.get("document_count", 0) > 0:
|
||||
await tool_ctx.info(
|
||||
"Project already indexed, use --force to re-index"
|
||||
)
|
||||
if show_stats:
|
||||
return self._format_stats(
|
||||
stats, abs_path, time.time() - start_time
|
||||
)
|
||||
return "Project is already indexed. Use --force to re-index."
|
||||
|
||||
# Prepare file patterns
|
||||
if file_patterns is None:
|
||||
# Default patterns for code files
|
||||
file_patterns = [
|
||||
"*.py",
|
||||
"*.js",
|
||||
"*.ts",
|
||||
"*.jsx",
|
||||
"*.tsx",
|
||||
"*.java",
|
||||
"*.cpp",
|
||||
"*.c",
|
||||
"*.h",
|
||||
"*.hpp",
|
||||
"*.go",
|
||||
"*.rs",
|
||||
"*.rb",
|
||||
"*.php",
|
||||
"*.swift",
|
||||
"*.kt",
|
||||
"*.scala",
|
||||
"*.cs",
|
||||
"*.vb",
|
||||
"*.fs",
|
||||
"*.sh",
|
||||
"*.bash",
|
||||
"*.zsh",
|
||||
"*.fish",
|
||||
"*.md",
|
||||
"*.rst",
|
||||
"*.txt",
|
||||
"*.json",
|
||||
"*.yaml",
|
||||
"*.yml",
|
||||
"*.toml",
|
||||
"*.ini",
|
||||
"*.cfg",
|
||||
"*.conf",
|
||||
"*.html",
|
||||
"*.css",
|
||||
"*.scss",
|
||||
"*.sass",
|
||||
"*.less",
|
||||
"*.sql",
|
||||
"*.graphql",
|
||||
"*.proto",
|
||||
"Dockerfile",
|
||||
"Makefile",
|
||||
"*.mk",
|
||||
".gitignore",
|
||||
".dockerignore",
|
||||
"requirements.txt",
|
||||
"package.json",
|
||||
"Cargo.toml",
|
||||
"go.mod",
|
||||
"pom.xml",
|
||||
]
|
||||
|
||||
# Clear existing index if force
|
||||
if force:
|
||||
await tool_ctx.info("Clearing existing index...")
|
||||
await vector_store.clear()
|
||||
|
||||
# Index files
|
||||
await tool_ctx.info("Indexing files...")
|
||||
indexed_files = 0
|
||||
total_size = 0
|
||||
errors = []
|
||||
|
||||
for pattern in file_patterns:
|
||||
pattern_files = await self._find_files(abs_path, pattern)
|
||||
for file_path in pattern_files:
|
||||
try:
|
||||
# Check file size (skip very large files)
|
||||
file_size = os.path.getsize(file_path)
|
||||
if file_size > 10 * 1024 * 1024: # 10MB
|
||||
await tool_ctx.warning(f"Skipping large file: {file_path}")
|
||||
continue
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
except UnicodeDecodeError:
|
||||
# Skip binary files
|
||||
continue
|
||||
|
||||
# Index the file
|
||||
rel_path = os.path.relpath(file_path, abs_path)
|
||||
await vector_store.index_document(
|
||||
content=content,
|
||||
metadata={
|
||||
"type": "file",
|
||||
"path": rel_path,
|
||||
"absolute_path": file_path,
|
||||
"size": file_size,
|
||||
"extension": Path(file_path).suffix,
|
||||
},
|
||||
)
|
||||
indexed_files += 1
|
||||
total_size += file_size
|
||||
|
||||
if indexed_files % 100 == 0:
|
||||
await tool_ctx.info(f"Indexed {indexed_files} files...")
|
||||
|
||||
except Exception as e:
|
||||
errors.append(f"{file_path}: {str(e)}")
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Indexed {indexed_files} files ({total_size / 1024 / 1024:.1f} MB)"
|
||||
)
|
||||
|
||||
# Index git history if requested
|
||||
git_stats = {}
|
||||
if include_git_history and os.path.exists(os.path.join(abs_path, ".git")):
|
||||
await tool_ctx.info("Indexing git history...")
|
||||
|
||||
git_ingester = GitIngester(vector_store)
|
||||
git_stats = await git_ingester.ingest_repository(
|
||||
repo_path=abs_path,
|
||||
include_history=True,
|
||||
include_diffs=True,
|
||||
include_blame=True,
|
||||
file_patterns=file_patterns,
|
||||
)
|
||||
|
||||
await tool_ctx.info(
|
||||
f"Indexed {git_stats.get('commits_indexed', 0)} commits, {git_stats.get('diffs_indexed', 0)} diffs"
|
||||
)
|
||||
|
||||
# Get final statistics
|
||||
if show_stats:
|
||||
stats = await vector_store.get_stats()
|
||||
stats.update(
|
||||
{
|
||||
"files_indexed": indexed_files,
|
||||
"total_size_mb": total_size / 1024 / 1024,
|
||||
"errors": len(errors),
|
||||
**git_stats,
|
||||
}
|
||||
)
|
||||
result = self._format_stats(stats, abs_path, time.time() - start_time)
|
||||
|
||||
if errors:
|
||||
result += f"\n\nErrors ({len(errors)}):\n"
|
||||
result += "\n".join(errors[:10]) # Show first 10 errors
|
||||
if len(errors) > 10:
|
||||
result += f"\n... and {len(errors) - 10} more errors"
|
||||
|
||||
return result
|
||||
else:
|
||||
return f"Successfully indexed {indexed_files} files"
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Indexing failed: {str(e)}")
|
||||
return f"Error during indexing: {str(e)}"
|
||||
|
||||
async def _find_files(self, base_path: str, pattern: str) -> list[str]:
|
||||
"""Find files matching a pattern.
|
||||
|
||||
Args:
|
||||
base_path: Base directory to search
|
||||
pattern: File pattern to match
|
||||
|
||||
Returns:
|
||||
List of matching file paths
|
||||
"""
|
||||
import glob
|
||||
|
||||
# Use glob to find files
|
||||
if pattern.startswith("*."):
|
||||
# Extension pattern
|
||||
files = glob.glob(
|
||||
os.path.join(base_path, "**", pattern),
|
||||
recursive=True,
|
||||
)
|
||||
else:
|
||||
# Exact filename
|
||||
files = glob.glob(
|
||||
os.path.join(base_path, "**", pattern),
|
||||
recursive=True,
|
||||
)
|
||||
|
||||
# Filter out hidden directories and common ignore patterns
|
||||
filtered_files = []
|
||||
ignore_dirs = {
|
||||
".git",
|
||||
"__pycache__",
|
||||
"node_modules",
|
||||
".venv",
|
||||
"venv",
|
||||
"dist",
|
||||
"build",
|
||||
}
|
||||
|
||||
for file_path in files:
|
||||
# Check if any parent directory is in ignore list
|
||||
parts = Path(file_path).parts
|
||||
if any(part in ignore_dirs for part in parts):
|
||||
continue
|
||||
if any(part.startswith(".") and part != "." for part in parts[:-1]):
|
||||
continue # Skip hidden directories (but allow hidden files like .gitignore)
|
||||
filtered_files.append(file_path)
|
||||
|
||||
return filtered_files
|
||||
|
||||
def _format_stats(self, stats: dict, path: str, elapsed_time: float) -> str:
|
||||
"""Format statistics for display.
|
||||
|
||||
Args:
|
||||
stats: Statistics dictionary
|
||||
path: Indexed path
|
||||
elapsed_time: Time taken for indexing
|
||||
|
||||
Returns:
|
||||
Formatted statistics string
|
||||
"""
|
||||
result = f"=== Index Statistics for {path} ===\n\n"
|
||||
|
||||
# Basic stats
|
||||
result += f"Indexing completed in {elapsed_time:.1f} seconds\n\n"
|
||||
|
||||
result += "Content Statistics:\n"
|
||||
result += f" Documents: {stats.get('document_count', 0):,}\n"
|
||||
result += f" Files indexed: {stats.get('files_indexed', 0):,}\n"
|
||||
result += f" Total size: {stats.get('total_size_mb', 0):.1f} MB\n"
|
||||
|
||||
if stats.get("commits_indexed", 0) > 0:
|
||||
result += f"\nGit History:\n"
|
||||
result += f" Commits: {stats.get('commits_indexed', 0):,}\n"
|
||||
result += f" Diffs: {stats.get('diffs_indexed', 0):,}\n"
|
||||
result += f" Blame entries: {stats.get('blame_entries', 0):,}\n"
|
||||
|
||||
# Vector store info
|
||||
result += f"\nVector Store:\n"
|
||||
result += f" Database: {stats.get('database_name', 'default')}\n"
|
||||
result += f" Table: {stats.get('table_name', 'documents')}\n"
|
||||
result += f" Vectors: {stats.get('vector_count', stats.get('document_count', 0)):,}\n"
|
||||
|
||||
if stats.get("errors", 0) > 0:
|
||||
result += f"\nErrors: {stats.get('errors', 0)}\n"
|
||||
|
||||
return result
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
# Tool registration is handled by the ToolRegistry
|
||||
pass
|
||||
@@ -1,902 +0,0 @@
|
||||
"""Infinity vector database integration for Hanzo AI."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import hashlib
|
||||
from typing import Any, Dict, List, Tuple, Optional
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass
|
||||
|
||||
try:
|
||||
import infinity_embedded
|
||||
|
||||
INFINITY_AVAILABLE = True
|
||||
except ImportError:
|
||||
# No silent mock: the random-vector mock makes search meaningless, so it is
|
||||
# OPT-IN only (HANZO_VECTOR_ALLOW_MOCK=1) for tests. Otherwise the local store
|
||||
# is simply unavailable — callers should use the cloud VectorTool (default).
|
||||
infinity_embedded = None
|
||||
INFINITY_AVAILABLE = False
|
||||
if os.environ.get("HANZO_VECTOR_ALLOW_MOCK") == "1":
|
||||
from . import mock_infinity as infinity_embedded
|
||||
|
||||
INFINITY_AVAILABLE = True
|
||||
|
||||
from .ast_analyzer import Symbol, FileAST, ASTAnalyzer, create_symbol_embedding_text
|
||||
|
||||
|
||||
@dataclass
|
||||
class Document:
|
||||
"""Document representation for vector storage."""
|
||||
|
||||
id: str
|
||||
content: str
|
||||
metadata: Dict[str, Any]
|
||||
file_path: Optional[str] = None
|
||||
chunk_index: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
"""Search result from vector database."""
|
||||
|
||||
document: Document
|
||||
score: float
|
||||
distance: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class SymbolSearchResult:
|
||||
"""Search result for symbols."""
|
||||
|
||||
symbol: Symbol
|
||||
score: float
|
||||
context_document: Optional[Document] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class UnifiedSearchResult:
|
||||
"""Search result combining text, vector, and symbol search."""
|
||||
|
||||
type: str # 'document', 'symbol', 'reference'
|
||||
content: str
|
||||
file_path: str
|
||||
line_start: int
|
||||
line_end: int
|
||||
score: float
|
||||
search_type: str # 'text', 'vector', 'symbol', 'ast'
|
||||
metadata: Dict[str, Any]
|
||||
|
||||
|
||||
class InfinityVectorStore:
|
||||
"""Local vector database using Infinity."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_path: Optional[str] = None,
|
||||
embedding_model: str = "text-embedding-3-small",
|
||||
dimension: int = 1536, # Default for OpenAI text-embedding-3-small
|
||||
):
|
||||
"""Initialize the Infinity vector store.
|
||||
|
||||
Args:
|
||||
data_path: Path to store vector database (default: ~/.config/hanzo/vector-store)
|
||||
embedding_model: Embedding model to use
|
||||
dimension: Vector dimension (must match embedding model)
|
||||
"""
|
||||
if not INFINITY_AVAILABLE:
|
||||
raise ImportError(
|
||||
"infinity_embedded is required for vector store functionality"
|
||||
)
|
||||
|
||||
# Set up data path
|
||||
if data_path:
|
||||
self.data_path = Path(data_path)
|
||||
else:
|
||||
from hanzo_mcp.config.settings import get_config_dir
|
||||
|
||||
self.data_path = get_config_dir() / "vector-store"
|
||||
|
||||
self.data_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.embedding_model = embedding_model
|
||||
self.dimension = dimension
|
||||
|
||||
# Initialize AST analyzer
|
||||
self.ast_analyzer = ASTAnalyzer()
|
||||
|
||||
# Connect to Infinity
|
||||
self.infinity = infinity_embedded.connect(str(self.data_path))
|
||||
self.db = self.infinity.get_database("hanzo_mcp")
|
||||
|
||||
# Initialize tables
|
||||
self._initialize_tables()
|
||||
|
||||
def _initialize_tables(self):
|
||||
"""Initialize database tables if they don't exist."""
|
||||
# Documents table
|
||||
try:
|
||||
self.documents_table = self.db.get_table("documents")
|
||||
except Exception:
|
||||
self.documents_table = self.db.create_table(
|
||||
"documents",
|
||||
{
|
||||
"id": {"type": "varchar"},
|
||||
"content": {"type": "varchar"},
|
||||
"file_path": {"type": "varchar"},
|
||||
"chunk_index": {"type": "integer"},
|
||||
"metadata": {"type": "varchar"}, # JSON string
|
||||
"embedding": {"type": f"vector,{self.dimension},float"},
|
||||
},
|
||||
)
|
||||
|
||||
# Symbols table for code symbols
|
||||
try:
|
||||
self.symbols_table = self.db.get_table("symbols")
|
||||
except Exception:
|
||||
self.symbols_table = self.db.create_table(
|
||||
"symbols",
|
||||
{
|
||||
"id": {"type": "varchar"},
|
||||
"name": {"type": "varchar"},
|
||||
"type": {"type": "varchar"}, # function, class, variable, etc.
|
||||
"file_path": {"type": "varchar"},
|
||||
"line_start": {"type": "integer"},
|
||||
"line_end": {"type": "integer"},
|
||||
"scope": {"type": "varchar"},
|
||||
"parent": {"type": "varchar"},
|
||||
"signature": {"type": "varchar"},
|
||||
"docstring": {"type": "varchar"},
|
||||
"metadata": {"type": "varchar"}, # JSON string
|
||||
"embedding": {"type": f"vector,{self.dimension},float"},
|
||||
},
|
||||
)
|
||||
|
||||
# AST table for storing complete file ASTs
|
||||
try:
|
||||
self.ast_table = self.db.get_table("ast_files")
|
||||
except Exception:
|
||||
self.ast_table = self.db.create_table(
|
||||
"ast_files",
|
||||
{
|
||||
"file_path": {"type": "varchar"},
|
||||
"file_hash": {"type": "varchar"},
|
||||
"language": {"type": "varchar"},
|
||||
"ast_data": {"type": "varchar"}, # JSON string of complete AST
|
||||
"last_updated": {"type": "varchar"}, # ISO timestamp
|
||||
},
|
||||
)
|
||||
|
||||
# References table for cross-file references
|
||||
try:
|
||||
self.references_table = self.db.get_table("references")
|
||||
except Exception:
|
||||
self.references_table = self.db.create_table(
|
||||
"references",
|
||||
{
|
||||
"id": {"type": "varchar"},
|
||||
"source_file": {"type": "varchar"},
|
||||
"target_file": {"type": "varchar"},
|
||||
"symbol_name": {"type": "varchar"},
|
||||
"reference_type": {
|
||||
"type": "varchar"
|
||||
}, # import, call, inheritance, etc.
|
||||
"line_number": {"type": "integer"},
|
||||
"metadata": {"type": "varchar"}, # JSON string
|
||||
},
|
||||
)
|
||||
|
||||
def _generate_doc_id(
|
||||
self, content: str, file_path: str = "", chunk_index: int = 0
|
||||
) -> str:
|
||||
"""Generate a unique document ID."""
|
||||
content_hash = hashlib.sha256(content.encode()).hexdigest()[:16]
|
||||
path_hash = hashlib.sha256(file_path.encode()).hexdigest()[:8]
|
||||
return f"doc_{path_hash}_{chunk_index}_{content_hash}"
|
||||
|
||||
def add_document(
|
||||
self,
|
||||
content: str,
|
||||
metadata: Dict[str, Any] = None,
|
||||
file_path: Optional[str] = None,
|
||||
chunk_index: int = 0,
|
||||
embedding: Optional[List[float]] = None,
|
||||
) -> str:
|
||||
"""Add a document to the vector store.
|
||||
|
||||
Args:
|
||||
content: Document content
|
||||
metadata: Additional metadata
|
||||
file_path: Source file path
|
||||
chunk_index: Chunk index if document is part of larger file
|
||||
embedding: Pre-computed embedding (if None, will compute)
|
||||
|
||||
Returns:
|
||||
Document ID
|
||||
"""
|
||||
doc_id = self._generate_doc_id(content, file_path or "", chunk_index)
|
||||
|
||||
# Generate embedding if not provided
|
||||
if embedding is None:
|
||||
embedding = self._generate_embedding(content)
|
||||
|
||||
# Prepare metadata
|
||||
metadata = metadata or {}
|
||||
metadata_json = json.dumps(metadata)
|
||||
|
||||
# Insert document
|
||||
self.documents_table.insert(
|
||||
[
|
||||
{
|
||||
"id": doc_id,
|
||||
"content": content,
|
||||
"file_path": file_path or "",
|
||||
"chunk_index": chunk_index,
|
||||
"metadata": metadata_json,
|
||||
"embedding": embedding,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
return doc_id
|
||||
|
||||
def add_file(
|
||||
self,
|
||||
file_path: str,
|
||||
chunk_size: int = 1000,
|
||||
chunk_overlap: int = 200,
|
||||
metadata: Dict[str, Any] = None,
|
||||
) -> List[str]:
|
||||
"""Add a file to the vector store by chunking it.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file to add
|
||||
chunk_size: Maximum characters per chunk
|
||||
chunk_overlap: Characters to overlap between chunks
|
||||
metadata: Additional metadata for all chunks
|
||||
|
||||
Returns:
|
||||
List of document IDs for all chunks
|
||||
"""
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
# Read file content
|
||||
try:
|
||||
content = path.read_text(encoding="utf-8")
|
||||
except UnicodeDecodeError:
|
||||
# Try with different encoding
|
||||
content = path.read_text(encoding="latin-1")
|
||||
|
||||
# Chunk the content
|
||||
chunks = self._chunk_text(content, chunk_size, chunk_overlap)
|
||||
|
||||
# Add metadata
|
||||
file_metadata = metadata or {}
|
||||
file_metadata.update(
|
||||
{
|
||||
"file_name": path.name,
|
||||
"file_extension": path.suffix,
|
||||
"file_size": path.stat().st_size,
|
||||
}
|
||||
)
|
||||
|
||||
# Add each chunk
|
||||
doc_ids = []
|
||||
for i, chunk in enumerate(chunks):
|
||||
chunk_metadata = file_metadata.copy()
|
||||
chunk_metadata["chunk_number"] = i
|
||||
chunk_metadata["total_chunks"] = len(chunks)
|
||||
|
||||
doc_id = self.add_document(
|
||||
content=chunk,
|
||||
metadata=chunk_metadata,
|
||||
file_path=str(path),
|
||||
chunk_index=i,
|
||||
)
|
||||
doc_ids.append(doc_id)
|
||||
|
||||
return doc_ids
|
||||
|
||||
def add_file_with_ast(
|
||||
self,
|
||||
file_path: str,
|
||||
chunk_size: int = 1000,
|
||||
chunk_overlap: int = 200,
|
||||
metadata: Dict[str, Any] = None,
|
||||
) -> Tuple[List[str], Optional[FileAST]]:
|
||||
"""Add a file with full AST analysis and symbol extraction.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file to add
|
||||
chunk_size: Maximum characters per chunk for content
|
||||
chunk_overlap: Characters to overlap between chunks
|
||||
metadata: Additional metadata for all chunks
|
||||
|
||||
Returns:
|
||||
Tuple of (document IDs for content chunks, FileAST object)
|
||||
"""
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
# First add file content using existing method
|
||||
doc_ids = self.add_file(file_path, chunk_size, chunk_overlap, metadata)
|
||||
|
||||
# Analyze AST and symbols
|
||||
file_ast = self.ast_analyzer.analyze_file(file_path)
|
||||
if not file_ast:
|
||||
return doc_ids, None
|
||||
|
||||
# Store complete AST
|
||||
self._store_file_ast(file_ast)
|
||||
|
||||
# Store individual symbols with embeddings
|
||||
self._store_symbols(file_ast.symbols)
|
||||
|
||||
# Store cross-references
|
||||
self._store_references(file_ast)
|
||||
|
||||
return doc_ids, file_ast
|
||||
|
||||
def _store_file_ast(self, file_ast: FileAST):
|
||||
"""Store complete file AST information."""
|
||||
from datetime import datetime
|
||||
|
||||
# Remove existing AST for this file
|
||||
try:
|
||||
self.ast_table.delete(f"file_path = '{file_ast.file_path}'")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Insert new AST
|
||||
self.ast_table.insert(
|
||||
[
|
||||
{
|
||||
"file_path": file_ast.file_path,
|
||||
"file_hash": file_ast.file_hash,
|
||||
"language": file_ast.language,
|
||||
"ast_data": json.dumps(file_ast.to_dict()),
|
||||
"last_updated": datetime.now().isoformat(),
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
def _store_symbols(self, symbols: List[Symbol]):
|
||||
"""Store symbols with vector embeddings."""
|
||||
if not symbols:
|
||||
return
|
||||
|
||||
# Remove existing symbols for these files
|
||||
file_paths = list(set(symbol.file_path for symbol in symbols))
|
||||
for file_path in file_paths:
|
||||
try:
|
||||
self.symbols_table.delete(f"file_path = '{file_path}'")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Insert new symbols
|
||||
symbol_records = []
|
||||
for symbol in symbols:
|
||||
# Create embedding text for symbol
|
||||
embedding_text = create_symbol_embedding_text(symbol)
|
||||
embedding = self._generate_embedding(embedding_text)
|
||||
|
||||
# Generate symbol ID
|
||||
symbol_id = self._generate_symbol_id(symbol)
|
||||
|
||||
# Prepare metadata
|
||||
symbol_metadata = {
|
||||
"references": symbol.references,
|
||||
"embedding_text": embedding_text,
|
||||
}
|
||||
|
||||
symbol_records.append(
|
||||
{
|
||||
"id": symbol_id,
|
||||
"name": symbol.name,
|
||||
"type": symbol.type,
|
||||
"file_path": symbol.file_path,
|
||||
"line_start": symbol.line_start,
|
||||
"line_end": symbol.line_end,
|
||||
"scope": symbol.scope or "",
|
||||
"parent": symbol.parent or "",
|
||||
"signature": symbol.signature or "",
|
||||
"docstring": symbol.docstring or "",
|
||||
"metadata": json.dumps(symbol_metadata),
|
||||
"embedding": embedding,
|
||||
}
|
||||
)
|
||||
|
||||
if symbol_records:
|
||||
self.symbols_table.insert(symbol_records)
|
||||
|
||||
def _store_references(self, file_ast: FileAST):
|
||||
"""Store cross-file references."""
|
||||
if not file_ast.dependencies:
|
||||
return
|
||||
|
||||
# Remove existing references for this file
|
||||
try:
|
||||
self.references_table.delete(f"source_file = '{file_ast.file_path}'")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Insert new references
|
||||
reference_records = []
|
||||
for i, dependency in enumerate(file_ast.dependencies):
|
||||
ref_id = f"{file_ast.file_path}_{dependency}_{i}"
|
||||
reference_records.append(
|
||||
{
|
||||
"id": ref_id,
|
||||
"source_file": file_ast.file_path,
|
||||
"target_file": dependency,
|
||||
"symbol_name": dependency,
|
||||
"reference_type": "import",
|
||||
"line_number": 0, # Could be enhanced to track actual line numbers
|
||||
"metadata": json.dumps({}),
|
||||
}
|
||||
)
|
||||
|
||||
if reference_records:
|
||||
self.references_table.insert(reference_records)
|
||||
|
||||
def _generate_symbol_id(self, symbol: Symbol) -> str:
|
||||
"""Generate unique symbol ID."""
|
||||
text = f"{symbol.file_path}_{symbol.type}_{symbol.name}_{symbol.line_start}"
|
||||
return hashlib.sha256(text.encode()).hexdigest()[:16]
|
||||
|
||||
def search_symbols(
|
||||
self,
|
||||
query: str,
|
||||
symbol_type: Optional[str] = None,
|
||||
file_path: Optional[str] = None,
|
||||
limit: int = 10,
|
||||
score_threshold: float = 0.0,
|
||||
) -> List[SymbolSearchResult]:
|
||||
"""Search for symbols using vector similarity.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
symbol_type: Filter by symbol type (function, class, variable, etc.)
|
||||
file_path: Filter by file path
|
||||
limit: Maximum number of results
|
||||
score_threshold: Minimum similarity score
|
||||
|
||||
Returns:
|
||||
List of symbol search results
|
||||
"""
|
||||
# Generate query embedding
|
||||
query_embedding = self._generate_embedding(query)
|
||||
|
||||
# Build search query
|
||||
search_query = self.symbols_table.output(["*"]).match_dense(
|
||||
"embedding",
|
||||
query_embedding,
|
||||
"float",
|
||||
"ip", # Inner product
|
||||
limit * 2, # Get more results for filtering
|
||||
)
|
||||
|
||||
# Apply filters
|
||||
if symbol_type:
|
||||
search_query = search_query.filter(f"type = '{symbol_type}'")
|
||||
if file_path:
|
||||
search_query = search_query.filter(f"file_path = '{file_path}'")
|
||||
|
||||
search_results = search_query.to_pl()
|
||||
|
||||
# Convert to SymbolSearchResult objects
|
||||
results = []
|
||||
for row in search_results.iter_rows(named=True):
|
||||
score = row.get("score", 0.0)
|
||||
if score >= score_threshold:
|
||||
# Parse metadata
|
||||
try:
|
||||
metadata = json.loads(row["metadata"])
|
||||
except Exception:
|
||||
metadata = {}
|
||||
|
||||
# Create Symbol object
|
||||
symbol = Symbol(
|
||||
name=row["name"],
|
||||
type=row["type"],
|
||||
file_path=row["file_path"],
|
||||
line_start=row["line_start"],
|
||||
line_end=row["line_end"],
|
||||
column_start=0, # Not stored in table
|
||||
column_end=0, # Not stored in table
|
||||
scope=row["scope"],
|
||||
parent=row["parent"] if row["parent"] else None,
|
||||
docstring=row["docstring"] if row["docstring"] else None,
|
||||
signature=row["signature"] if row["signature"] else None,
|
||||
references=metadata.get("references", []),
|
||||
)
|
||||
|
||||
results.append(
|
||||
SymbolSearchResult(
|
||||
symbol=symbol,
|
||||
score=score,
|
||||
)
|
||||
)
|
||||
|
||||
return results[:limit]
|
||||
|
||||
def search_ast_nodes(
|
||||
self,
|
||||
file_path: str,
|
||||
node_type: Optional[str] = None,
|
||||
node_name: Optional[str] = None,
|
||||
) -> Optional[FileAST]:
|
||||
"""Search AST nodes within a specific file.
|
||||
|
||||
Args:
|
||||
file_path: File to search in
|
||||
node_type: Filter by AST node type
|
||||
node_name: Filter by node name
|
||||
|
||||
Returns:
|
||||
FileAST object if file found, None otherwise
|
||||
"""
|
||||
try:
|
||||
results = (
|
||||
self.ast_table.output(["*"])
|
||||
.filter(f"file_path = '{file_path}'")
|
||||
.to_pl()
|
||||
)
|
||||
|
||||
if len(results) == 0:
|
||||
return None
|
||||
|
||||
row = next(results.iter_rows(named=True))
|
||||
ast_data = json.loads(row["ast_data"])
|
||||
|
||||
# Reconstruct FileAST object
|
||||
file_ast = FileAST(
|
||||
file_path=ast_data["file_path"],
|
||||
file_hash=ast_data["file_hash"],
|
||||
language=ast_data["language"],
|
||||
symbols=[Symbol(**s) for s in ast_data["symbols"]],
|
||||
ast_nodes=[], # Would need custom deserialization for ASTNode
|
||||
imports=ast_data["imports"],
|
||||
exports=ast_data["exports"],
|
||||
dependencies=ast_data["dependencies"],
|
||||
)
|
||||
|
||||
return file_ast
|
||||
|
||||
except Exception as e:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.error(f"Error searching AST nodes: {e}")
|
||||
return None
|
||||
|
||||
def get_file_references(self, file_path: str) -> List[Dict[str, Any]]:
|
||||
"""Get all files that reference the given file.
|
||||
|
||||
Args:
|
||||
file_path: File to find references for
|
||||
|
||||
Returns:
|
||||
List of reference information
|
||||
"""
|
||||
try:
|
||||
results = (
|
||||
self.references_table.output(["*"])
|
||||
.filter(f"target_file = '{file_path}'")
|
||||
.to_pl()
|
||||
)
|
||||
|
||||
references = []
|
||||
for row in results.iter_rows(named=True):
|
||||
references.append(
|
||||
{
|
||||
"source_file": row["source_file"],
|
||||
"symbol_name": row["symbol_name"],
|
||||
"reference_type": row["reference_type"],
|
||||
"line_number": row["line_number"],
|
||||
}
|
||||
)
|
||||
|
||||
return references
|
||||
|
||||
except Exception as e:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.error(f"Error getting file references: {e}")
|
||||
return []
|
||||
|
||||
def search(
|
||||
self,
|
||||
query: str,
|
||||
limit: int = 10,
|
||||
score_threshold: float = 0.0,
|
||||
filters: Dict[str, Any] = None,
|
||||
) -> List[SearchResult]:
|
||||
"""Search for similar documents.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
limit: Maximum number of results
|
||||
score_threshold: Minimum similarity score
|
||||
filters: Metadata filters (not yet implemented)
|
||||
|
||||
Returns:
|
||||
List of search results
|
||||
"""
|
||||
# Generate query embedding
|
||||
query_embedding = self._generate_embedding(query)
|
||||
|
||||
# Perform vector search
|
||||
search_results = (
|
||||
self.documents_table.output(["*"])
|
||||
.match_dense(
|
||||
"embedding",
|
||||
query_embedding,
|
||||
"float",
|
||||
"ip", # Inner product (cosine similarity)
|
||||
limit,
|
||||
)
|
||||
.to_pl()
|
||||
)
|
||||
|
||||
# Convert to SearchResult objects
|
||||
results = []
|
||||
for row in search_results.iter_rows(named=True):
|
||||
# Parse metadata
|
||||
try:
|
||||
metadata = json.loads(row["metadata"])
|
||||
except Exception:
|
||||
metadata = {}
|
||||
|
||||
# Create document
|
||||
document = Document(
|
||||
id=row["id"],
|
||||
content=row["content"],
|
||||
metadata=metadata,
|
||||
file_path=row["file_path"] if row["file_path"] else None,
|
||||
chunk_index=row["chunk_index"],
|
||||
)
|
||||
|
||||
# Score is the similarity (higher is better)
|
||||
score = row.get("score", 0.0)
|
||||
distance = 1.0 - score # Convert similarity to distance
|
||||
|
||||
if score >= score_threshold:
|
||||
results.append(
|
||||
SearchResult(
|
||||
document=document,
|
||||
score=score,
|
||||
distance=distance,
|
||||
)
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
def delete_document(self, doc_id: str) -> bool:
|
||||
"""Delete a document by ID.
|
||||
|
||||
Args:
|
||||
doc_id: Document ID to delete
|
||||
|
||||
Returns:
|
||||
True if document was deleted
|
||||
"""
|
||||
try:
|
||||
self.documents_table.delete(f"id = '{doc_id}'")
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def delete_file(self, file_path: str) -> int:
|
||||
"""Delete all documents from a specific file.
|
||||
|
||||
Args:
|
||||
file_path: File path to delete documents for
|
||||
|
||||
Returns:
|
||||
Number of documents deleted
|
||||
"""
|
||||
try:
|
||||
# Get count first
|
||||
results = (
|
||||
self.documents_table.output(["id"])
|
||||
.filter(f"file_path = '{file_path}'")
|
||||
.to_pl()
|
||||
)
|
||||
count = len(results)
|
||||
|
||||
# Delete all documents for this file
|
||||
self.documents_table.delete(f"file_path = '{file_path}'")
|
||||
return count
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
def list_files(self) -> List[Dict[str, Any]]:
|
||||
"""List all indexed files.
|
||||
|
||||
Returns:
|
||||
List of file information
|
||||
"""
|
||||
try:
|
||||
results = self.documents_table.output(["file_path", "metadata"]).to_pl()
|
||||
|
||||
files = {}
|
||||
for row in results.iter_rows(named=True):
|
||||
file_path = row["file_path"]
|
||||
if file_path and file_path not in files:
|
||||
try:
|
||||
metadata = json.loads(row["metadata"])
|
||||
files[file_path] = {
|
||||
"file_path": file_path,
|
||||
"file_name": metadata.get(
|
||||
"file_name", Path(file_path).name
|
||||
),
|
||||
"file_size": metadata.get("file_size", 0),
|
||||
"total_chunks": metadata.get("total_chunks", 1),
|
||||
}
|
||||
except Exception:
|
||||
files[file_path] = {
|
||||
"file_path": file_path,
|
||||
"file_name": Path(file_path).name,
|
||||
}
|
||||
|
||||
return list(files.values())
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
def _chunk_text(self, text: str, chunk_size: int, overlap: int) -> List[str]:
|
||||
"""Split text into overlapping chunks."""
|
||||
if len(text) <= chunk_size:
|
||||
return [text]
|
||||
|
||||
chunks = []
|
||||
start = 0
|
||||
|
||||
while start < len(text):
|
||||
end = start + chunk_size
|
||||
|
||||
# Try to break at word boundary
|
||||
if end < len(text):
|
||||
# Look back for a good break point
|
||||
break_point = end
|
||||
for i in range(end - 100, start + 100, -1):
|
||||
if i > 0 and text[i] in "\n\r.!?":
|
||||
break_point = i + 1
|
||||
break
|
||||
end = break_point
|
||||
|
||||
chunk = text[start:end].strip()
|
||||
if chunk:
|
||||
chunks.append(chunk)
|
||||
|
||||
start = max(start + chunk_size - overlap, end)
|
||||
|
||||
return chunks
|
||||
|
||||
def _generate_embedding(self, text: str) -> List[float]:
|
||||
"""Generate an embedding for text.
|
||||
|
||||
The local Infinity store ships no embedder. Rather than return random
|
||||
vectors (which make similarity search meaningless), this raises unless
|
||||
the mock is explicitly enabled (HANZO_VECTOR_ALLOW_MOCK=1). For real
|
||||
vectors, use the cloud VectorTool (api.hanzo.ai /v1/embeddings).
|
||||
"""
|
||||
if os.environ.get("HANZO_VECTOR_ALLOW_MOCK") == "1":
|
||||
import random
|
||||
|
||||
return [random.random() for _ in range(self.dimension)]
|
||||
raise NotImplementedError(
|
||||
"Local Infinity store has no embedder; use the cloud VectorTool "
|
||||
"(default) for real embeddings, or set HANZO_VECTOR_ALLOW_MOCK=1 "
|
||||
"for random test vectors."
|
||||
)
|
||||
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get statistics about the vector store.
|
||||
|
||||
Returns:
|
||||
Dictionary with statistics
|
||||
"""
|
||||
try:
|
||||
# Get document count
|
||||
doc_count_result = self.documents_table.output(["count(*)"]).to_pl()
|
||||
doc_count = doc_count_result.item(0, 0) if len(doc_count_result) > 0 else 0
|
||||
|
||||
# Get unique file count
|
||||
file_result = self.documents_table.output(["file_path"]).to_pl()
|
||||
unique_files = set()
|
||||
for row in file_result.iter_rows():
|
||||
if row[0]:
|
||||
unique_files.add(row[0])
|
||||
|
||||
# Get symbol count
|
||||
symbol_count = 0
|
||||
try:
|
||||
symbol_result = self.symbols_table.output(["count(*)"]).to_pl()
|
||||
symbol_count = symbol_result.item(0, 0) if len(symbol_result) > 0 else 0
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Get AST count
|
||||
ast_count = 0
|
||||
try:
|
||||
ast_result = self.ast_table.output(["count(*)"]).to_pl()
|
||||
ast_count = ast_result.item(0, 0) if len(ast_result) > 0 else 0
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return {
|
||||
"document_count": doc_count,
|
||||
"vector_count": doc_count, # Each document has a vector
|
||||
"unique_files": len(unique_files),
|
||||
"symbol_count": symbol_count,
|
||||
"ast_count": ast_count,
|
||||
"database_name": self.db_name,
|
||||
"table_name": "documents",
|
||||
"dimension": self.dimension,
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
"error": str(e),
|
||||
"document_count": 0,
|
||||
"vector_count": 0,
|
||||
}
|
||||
|
||||
async def clear(self) -> bool:
|
||||
"""Clear all data from the vector store.
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
try:
|
||||
# Delete all records from all tables
|
||||
self.documents_table.delete()
|
||||
|
||||
try:
|
||||
self.symbols_table.delete()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
self.ast_table.delete()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
self.references_table.delete()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.error(f"Error clearing vector store: {e}")
|
||||
return False
|
||||
|
||||
async def index_document(
|
||||
self,
|
||||
content: str,
|
||||
metadata: Dict[str, Any] = None,
|
||||
) -> str:
|
||||
"""Async version of add_document for consistency.
|
||||
|
||||
Args:
|
||||
content: Document content
|
||||
metadata: Additional metadata
|
||||
|
||||
Returns:
|
||||
Document ID
|
||||
"""
|
||||
file_path = metadata.get("path") if metadata else None
|
||||
return self.add_document(content, metadata, file_path)
|
||||
|
||||
def close(self):
|
||||
"""Close the database connection."""
|
||||
if hasattr(self, "infinity"):
|
||||
self.infinity.disconnect()
|
||||
@@ -1,161 +0,0 @@
|
||||
"""Mock implementation of infinity_embedded for testing on unsupported platforms."""
|
||||
|
||||
import random
|
||||
from typing import Any, Dict, List
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class MockTable:
|
||||
"""Mock implementation of an Infinity table."""
|
||||
|
||||
def __init__(self, name: str, schema: Dict[str, Any]):
|
||||
self.name = name
|
||||
self.schema = schema
|
||||
self.data = []
|
||||
self._id_counter = 0
|
||||
|
||||
def insert(self, records: List[Dict[str, Any]]):
|
||||
"""Insert records into the table."""
|
||||
for record in records:
|
||||
# Add an internal ID if not present
|
||||
if "id" not in record:
|
||||
record["_internal_id"] = self._id_counter
|
||||
self._id_counter += 1
|
||||
self.data.append(record)
|
||||
|
||||
def delete(self, condition: str):
|
||||
"""Delete records matching condition."""
|
||||
# Simple implementation - just clear for now
|
||||
self.data = [r for r in self.data if not self._eval_condition(r, condition)]
|
||||
|
||||
def output(self, columns: List[str]):
|
||||
"""Start a query chain."""
|
||||
return MockQuery(self, columns)
|
||||
|
||||
def _eval_condition(self, record: Dict[str, Any], condition: str) -> bool:
|
||||
"""Evaluate a simple condition."""
|
||||
# Very basic implementation
|
||||
if "=" in condition:
|
||||
field, value = condition.split("=", 1)
|
||||
field = field.strip()
|
||||
value = value.strip().strip("'\"")
|
||||
return str(record.get(field, "")) == value
|
||||
return False
|
||||
|
||||
|
||||
class MockQuery:
|
||||
"""Mock query builder."""
|
||||
|
||||
def __init__(self, table: MockTable, columns: List[str]):
|
||||
self.table = table
|
||||
self.columns = columns
|
||||
self.filters = []
|
||||
self.vector_search = None
|
||||
self.limit_value = None
|
||||
|
||||
def filter(self, condition: str):
|
||||
"""Add a filter condition."""
|
||||
self.filters.append(condition)
|
||||
return self
|
||||
|
||||
def match_dense(
|
||||
self, column: str, vector: List[float], dtype: str, metric: str, limit: int
|
||||
):
|
||||
"""Add vector search."""
|
||||
self.vector_search = {
|
||||
"column": column,
|
||||
"vector": vector,
|
||||
"dtype": dtype,
|
||||
"metric": metric,
|
||||
"limit": limit,
|
||||
}
|
||||
self.limit_value = limit
|
||||
return self
|
||||
|
||||
def to_pl(self):
|
||||
"""Execute query and return polars-like result."""
|
||||
results = self.table.data.copy()
|
||||
|
||||
# Apply filters
|
||||
for condition in self.filters:
|
||||
results = [r for r in results if self.table._eval_condition(r, condition)]
|
||||
|
||||
# Apply vector search (mock similarity)
|
||||
if self.vector_search:
|
||||
# Add mock scores
|
||||
for r in results:
|
||||
r["score"] = random.uniform(0.5, 1.0)
|
||||
# Sort by score
|
||||
results.sort(key=lambda x: x.get("score", 0), reverse=True)
|
||||
# Limit results
|
||||
if self.limit_value:
|
||||
results = results[: self.limit_value]
|
||||
|
||||
# Return mock polars DataFrame
|
||||
return MockDataFrame(results)
|
||||
|
||||
|
||||
class MockDataFrame:
|
||||
"""Mock polars DataFrame."""
|
||||
|
||||
def __init__(self, data: List[Dict[str, Any]]):
|
||||
self.data = data
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
def iter_rows(self, named: bool = False):
|
||||
"""Iterate over rows."""
|
||||
if named:
|
||||
return iter(self.data)
|
||||
else:
|
||||
# Return tuples
|
||||
if not self.data:
|
||||
return iter([])
|
||||
keys = list(self.data[0].keys())
|
||||
return iter([tuple(row.get(k) for k in keys) for row in self.data])
|
||||
|
||||
|
||||
class MockDatabase:
|
||||
"""Mock implementation of an Infinity database."""
|
||||
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
self.tables = {}
|
||||
|
||||
def create_table(self, name: str, schema: Dict[str, Any]) -> MockTable:
|
||||
"""Create a new table."""
|
||||
table = MockTable(name, schema)
|
||||
self.tables[name] = table
|
||||
return table
|
||||
|
||||
def get_table(self, name: str) -> MockTable:
|
||||
"""Get an existing table."""
|
||||
if name not in self.tables:
|
||||
raise KeyError(f"Table {name} not found")
|
||||
return self.tables[name]
|
||||
|
||||
|
||||
class MockInfinity:
|
||||
"""Mock implementation of Infinity connection."""
|
||||
|
||||
def __init__(self, path: str):
|
||||
self.path = Path(path)
|
||||
self.databases = {}
|
||||
# Ensure directory exists
|
||||
self.path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def get_database(self, name: str) -> MockDatabase:
|
||||
"""Get or create a database."""
|
||||
if name not in self.databases:
|
||||
self.databases[name] = MockDatabase(name)
|
||||
return self.databases[name]
|
||||
|
||||
def disconnect(self):
|
||||
"""Disconnect from Infinity."""
|
||||
pass
|
||||
|
||||
|
||||
def connect(path: str) -> MockInfinity:
|
||||
"""Connect to Infinity (mock implementation)."""
|
||||
return MockInfinity(path)
|
||||
@@ -1,394 +0,0 @@
|
||||
"""Project-aware vector database management for Hanzo AI."""
|
||||
|
||||
import asyncio
|
||||
from typing import Any, Dict, List, Tuple, Optional
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from .index_config import IndexScope, IndexConfig
|
||||
from .infinity_store import SearchResult, InfinityVectorStore
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProjectInfo:
|
||||
"""Information about a detected project."""
|
||||
|
||||
root_path: Path
|
||||
llm_md_path: Path
|
||||
db_path: Path
|
||||
name: str
|
||||
|
||||
|
||||
class ProjectVectorManager:
|
||||
"""Manages project-aware vector databases."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
global_db_path: Optional[str] = None,
|
||||
embedding_model: str = "text-embedding-3-small",
|
||||
dimension: int = 1536,
|
||||
):
|
||||
"""Initialize the project vector manager.
|
||||
|
||||
Args:
|
||||
global_db_path: Path for global vector store (default: ~/.config/hanzo/db)
|
||||
embedding_model: Embedding model to use
|
||||
dimension: Vector dimension
|
||||
"""
|
||||
self.embedding_model = embedding_model
|
||||
self.dimension = dimension
|
||||
|
||||
# Set up index configuration
|
||||
self.index_config = IndexConfig()
|
||||
|
||||
# Set up global database path
|
||||
if global_db_path:
|
||||
self.global_db_path = Path(global_db_path)
|
||||
else:
|
||||
self.global_db_path = self.index_config.get_index_path("vector")
|
||||
|
||||
self.global_db_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Cache for project info and vector stores
|
||||
self.projects: Dict[str, ProjectInfo] = {}
|
||||
self.vector_stores: Dict[str, InfinityVectorStore] = {}
|
||||
self._global_store: Optional[InfinityVectorStore] = None
|
||||
|
||||
# Thread pool for parallel operations
|
||||
self.executor = ThreadPoolExecutor(max_workers=4)
|
||||
|
||||
def _get_global_store(self) -> InfinityVectorStore:
|
||||
"""Get or create the global vector store."""
|
||||
if self._global_store is None:
|
||||
self._global_store = InfinityVectorStore(
|
||||
data_path=str(self.global_db_path),
|
||||
embedding_model=self.embedding_model,
|
||||
dimension=self.dimension,
|
||||
)
|
||||
return self._global_store
|
||||
|
||||
def detect_projects(self, search_paths: List[str]) -> List[ProjectInfo]:
|
||||
"""Detect projects by finding LLM.md files.
|
||||
|
||||
Args:
|
||||
search_paths: List of paths to search for projects
|
||||
|
||||
Returns:
|
||||
List of detected project information
|
||||
"""
|
||||
projects = []
|
||||
|
||||
for search_path in search_paths:
|
||||
path = Path(search_path).resolve()
|
||||
|
||||
# Search for LLM.md files
|
||||
for llm_md_path in path.rglob("LLM.md"):
|
||||
project_root = llm_md_path.parent
|
||||
project_name = project_root.name
|
||||
|
||||
# Create .hanzo/db directory in project
|
||||
db_path = project_root / ".hanzo" / "db"
|
||||
db_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
project_info = ProjectInfo(
|
||||
root_path=project_root,
|
||||
llm_md_path=llm_md_path,
|
||||
db_path=db_path,
|
||||
name=project_name,
|
||||
)
|
||||
|
||||
projects.append(project_info)
|
||||
|
||||
# Cache project info
|
||||
project_key = str(project_root)
|
||||
self.projects[project_key] = project_info
|
||||
|
||||
return projects
|
||||
|
||||
def get_project_for_path(self, file_path: str) -> Optional[ProjectInfo]:
|
||||
"""Find the project that contains a given file path.
|
||||
|
||||
Args:
|
||||
file_path: File path to check
|
||||
|
||||
Returns:
|
||||
Project info if found, None otherwise
|
||||
"""
|
||||
path = Path(file_path).resolve()
|
||||
|
||||
# Check each known project
|
||||
for project_key, project_info in self.projects.items():
|
||||
try:
|
||||
# Check if path is within project root
|
||||
path.relative_to(project_info.root_path)
|
||||
return project_info
|
||||
except ValueError:
|
||||
# Path is not within this project
|
||||
continue
|
||||
|
||||
# Try to find project by walking up the directory tree
|
||||
current_path = path.parent if path.is_file() else path
|
||||
|
||||
while current_path != current_path.parent: # Stop at filesystem root
|
||||
llm_md_path = current_path / "LLM.md"
|
||||
if llm_md_path.exists():
|
||||
# Found a project, create and cache it
|
||||
db_path = current_path / ".hanzo" / "db"
|
||||
db_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
project_info = ProjectInfo(
|
||||
root_path=current_path,
|
||||
llm_md_path=llm_md_path,
|
||||
db_path=db_path,
|
||||
name=current_path.name,
|
||||
)
|
||||
|
||||
project_key = str(current_path)
|
||||
self.projects[project_key] = project_info
|
||||
return project_info
|
||||
|
||||
current_path = current_path.parent
|
||||
|
||||
return None
|
||||
|
||||
def get_vector_store(
|
||||
self, project_info: Optional[ProjectInfo] = None
|
||||
) -> InfinityVectorStore:
|
||||
"""Get vector store for a project or global store.
|
||||
|
||||
Args:
|
||||
project_info: Project to get store for, None for global store
|
||||
|
||||
Returns:
|
||||
Vector store instance
|
||||
"""
|
||||
# Check indexing scope
|
||||
if project_info:
|
||||
scope = self.index_config.get_scope(str(project_info.root_path))
|
||||
if scope == IndexScope.GLOBAL:
|
||||
# Even for project files, use global store if configured
|
||||
return self._get_global_store()
|
||||
else:
|
||||
return self._get_global_store()
|
||||
|
||||
# Use project-specific store
|
||||
project_key = str(project_info.root_path)
|
||||
|
||||
if project_key not in self.vector_stores:
|
||||
# Get index path based on configuration
|
||||
index_path = self.index_config.get_index_path(
|
||||
"vector", str(project_info.root_path)
|
||||
)
|
||||
index_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.vector_stores[project_key] = InfinityVectorStore(
|
||||
data_path=str(index_path),
|
||||
embedding_model=self.embedding_model,
|
||||
dimension=self.dimension,
|
||||
)
|
||||
|
||||
return self.vector_stores[project_key]
|
||||
|
||||
def add_file_to_appropriate_store(
|
||||
self,
|
||||
file_path: str,
|
||||
chunk_size: int = 1000,
|
||||
chunk_overlap: int = 200,
|
||||
metadata: Dict[str, Any] = None,
|
||||
) -> Tuple[List[str], Optional[ProjectInfo]]:
|
||||
"""Add a file to the appropriate vector store (project or global).
|
||||
|
||||
Args:
|
||||
file_path: Path to file to add
|
||||
chunk_size: Chunk size for text splitting
|
||||
chunk_overlap: Overlap between chunks
|
||||
metadata: Additional metadata
|
||||
|
||||
Returns:
|
||||
Tuple of (document IDs, project info or None for global)
|
||||
"""
|
||||
# Check if indexing is enabled
|
||||
if not self.index_config.is_indexing_enabled("vector"):
|
||||
return [], None
|
||||
|
||||
# Find project for this file
|
||||
project_info = self.get_project_for_path(file_path)
|
||||
|
||||
# Get appropriate vector store based on scope configuration
|
||||
vector_store = self.get_vector_store(project_info)
|
||||
|
||||
# Add file metadata
|
||||
file_metadata = metadata or {}
|
||||
if project_info:
|
||||
file_metadata["project_name"] = project_info.name
|
||||
file_metadata["project_root"] = str(project_info.root_path)
|
||||
# Check actual scope used
|
||||
scope = self.index_config.get_scope(str(project_info.root_path))
|
||||
file_metadata["index_scope"] = scope.value
|
||||
else:
|
||||
file_metadata["project_name"] = "global"
|
||||
file_metadata["index_scope"] = "global"
|
||||
|
||||
# Add file to store
|
||||
doc_ids = vector_store.add_file(
|
||||
file_path=file_path,
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
metadata=file_metadata,
|
||||
)
|
||||
|
||||
return doc_ids, project_info
|
||||
|
||||
async def search_all_projects(
|
||||
self,
|
||||
query: str,
|
||||
limit_per_project: int = 5,
|
||||
score_threshold: float = 0.0,
|
||||
include_global: bool = True,
|
||||
project_filter: Optional[List[str]] = None,
|
||||
) -> Dict[str, List[SearchResult]]:
|
||||
"""Search across all projects in parallel.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
limit_per_project: Maximum results per project
|
||||
score_threshold: Minimum similarity score
|
||||
include_global: Whether to include global store
|
||||
project_filter: List of project names to search (None for all)
|
||||
|
||||
Returns:
|
||||
Dictionary mapping project names to search results
|
||||
"""
|
||||
search_tasks = []
|
||||
project_names = []
|
||||
|
||||
# Add global store if requested
|
||||
if include_global:
|
||||
global_store = self._get_global_store()
|
||||
search_tasks.append(
|
||||
asyncio.get_event_loop().run_in_executor(
|
||||
self.executor,
|
||||
lambda: global_store.search(
|
||||
query, limit_per_project, score_threshold
|
||||
),
|
||||
)
|
||||
)
|
||||
project_names.append("global")
|
||||
|
||||
# Add project stores
|
||||
for _project_key, project_info in self.projects.items():
|
||||
# Apply project filter
|
||||
if project_filter and project_info.name not in project_filter:
|
||||
continue
|
||||
|
||||
vector_store = self.get_vector_store(project_info)
|
||||
search_tasks.append(
|
||||
asyncio.get_event_loop().run_in_executor(
|
||||
self.executor,
|
||||
lambda vs=vector_store: vs.search(
|
||||
query, limit_per_project, score_threshold
|
||||
),
|
||||
)
|
||||
)
|
||||
project_names.append(project_info.name)
|
||||
|
||||
# Execute all searches in parallel
|
||||
results = await asyncio.gather(*search_tasks, return_exceptions=True)
|
||||
|
||||
# Combine results
|
||||
combined_results = {}
|
||||
for i, result in enumerate(results):
|
||||
project_name = project_names[i]
|
||||
if isinstance(result, Exception):
|
||||
# Log error but continue
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.error(f"Error searching project {project_name}: {result}")
|
||||
combined_results[project_name] = []
|
||||
else:
|
||||
combined_results[project_name] = result
|
||||
|
||||
return combined_results
|
||||
|
||||
def search_project_by_path(
|
||||
self,
|
||||
file_path: str,
|
||||
query: str,
|
||||
limit: int = 10,
|
||||
score_threshold: float = 0.0,
|
||||
) -> List[SearchResult]:
|
||||
"""Search the project containing a specific file path.
|
||||
|
||||
Args:
|
||||
file_path: File path to determine project
|
||||
query: Search query
|
||||
limit: Maximum results
|
||||
score_threshold: Minimum similarity score
|
||||
|
||||
Returns:
|
||||
Search results from the appropriate project store
|
||||
"""
|
||||
project_info = self.get_project_for_path(file_path)
|
||||
vector_store = self.get_vector_store(project_info)
|
||||
|
||||
return vector_store.search(
|
||||
query=query,
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
|
||||
def get_project_stats(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get statistics for all projects.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping project names to stats
|
||||
"""
|
||||
stats = {}
|
||||
|
||||
# Global store stats
|
||||
try:
|
||||
global_store = self._get_global_store()
|
||||
global_files = global_store.list_files()
|
||||
stats["global"] = {
|
||||
"file_count": len(global_files),
|
||||
"db_path": str(self.global_db_path),
|
||||
}
|
||||
except Exception as e:
|
||||
stats["global"] = {"error": str(e)}
|
||||
|
||||
# Project store stats
|
||||
for _project_key, project_info in self.projects.items():
|
||||
try:
|
||||
vector_store = self.get_vector_store(project_info)
|
||||
project_files = vector_store.list_files()
|
||||
stats[project_info.name] = {
|
||||
"file_count": len(project_files),
|
||||
"db_path": str(project_info.db_path),
|
||||
"root_path": str(project_info.root_path),
|
||||
"llm_md_exists": project_info.llm_md_path.exists(),
|
||||
}
|
||||
except Exception as e:
|
||||
stats[project_info.name] = {"error": str(e)}
|
||||
|
||||
return stats
|
||||
|
||||
def cleanup(self):
|
||||
"""Close all vector stores and cleanup resources."""
|
||||
# Close all project stores
|
||||
for vector_store in self.vector_stores.values():
|
||||
try:
|
||||
vector_store.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Close global store
|
||||
if self._global_store:
|
||||
try:
|
||||
self._global_store.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Shutdown executor
|
||||
self.executor.shutdown(wait=False)
|
||||
@@ -1,329 +0,0 @@
|
||||
"""Unified vector store tool."""
|
||||
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
Unpack,
|
||||
Optional,
|
||||
Annotated,
|
||||
TypedDict,
|
||||
final,
|
||||
override,
|
||||
)
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import Field
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, PermissionManager, auto_timeout
|
||||
|
||||
from .project_manager import ProjectVectorManager
|
||||
|
||||
# Parameter types
|
||||
Action = Annotated[
|
||||
str,
|
||||
Field(
|
||||
description="Action: search (default), index, stats, clear",
|
||||
default="search",
|
||||
),
|
||||
]
|
||||
|
||||
Query = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Search query for semantic similarity",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Path = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="Path to index or search within",
|
||||
default=".",
|
||||
),
|
||||
]
|
||||
|
||||
Include = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="File pattern to include (e.g., '*.py')",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Exclude = Annotated[
|
||||
Optional[str],
|
||||
Field(
|
||||
description="File pattern to exclude",
|
||||
default=None,
|
||||
),
|
||||
]
|
||||
|
||||
Limit = Annotated[
|
||||
int,
|
||||
Field(
|
||||
description="Maximum results to return",
|
||||
default=10,
|
||||
),
|
||||
]
|
||||
|
||||
IncludeGit = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Include git history in indexing",
|
||||
default=True,
|
||||
),
|
||||
]
|
||||
|
||||
ForceReindex = Annotated[
|
||||
bool,
|
||||
Field(
|
||||
description="Force reindexing even if up to date",
|
||||
default=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class VectorParams(TypedDict, total=False):
|
||||
"""Parameters for vector tool."""
|
||||
|
||||
action: str
|
||||
query: Optional[str]
|
||||
path: Optional[str]
|
||||
include: Optional[str]
|
||||
exclude: Optional[str]
|
||||
limit: int
|
||||
include_git: bool
|
||||
force_reindex: bool
|
||||
|
||||
|
||||
@final
|
||||
class VectorTool(BaseTool):
|
||||
"""Unified vector store tool for semantic search."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
project_manager: ProjectVectorManager,
|
||||
):
|
||||
"""Initialize the vector tool."""
|
||||
super().__init__(permission_manager)
|
||||
self.project_manager = project_manager
|
||||
|
||||
@property
|
||||
@override
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "vector"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Semantic search with embeddings. Actions: search (default), index, stats, clear.
|
||||
|
||||
Usage:
|
||||
vector "find authentication logic"
|
||||
vector --action index --path ./src --include "*.py"
|
||||
vector --action stats
|
||||
vector --action clear --path ./old_code
|
||||
"""
|
||||
|
||||
@override
|
||||
@auto_timeout("vector")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[VectorParams],
|
||||
) -> str:
|
||||
"""Execute vector operation."""
|
||||
tool_ctx = self.create_tool_context(ctx)
|
||||
|
||||
# Extract action
|
||||
action = params.get("action", "search")
|
||||
|
||||
# Route to appropriate handler
|
||||
if action == "search":
|
||||
return await self._handle_search(params, tool_ctx)
|
||||
elif action == "index":
|
||||
return await self._handle_index(params, tool_ctx)
|
||||
elif action == "stats":
|
||||
return await self._handle_stats(params, tool_ctx)
|
||||
elif action == "clear":
|
||||
return await self._handle_clear(params, tool_ctx)
|
||||
else:
|
||||
return f"Error: Unknown action '{action}'. Valid actions: search, index, stats, clear"
|
||||
|
||||
async def _handle_search(self, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Handle semantic search."""
|
||||
query = params.get("query")
|
||||
if not query:
|
||||
return "Error: query is required for search action"
|
||||
|
||||
path = params.get("path", ".")
|
||||
limit = params.get("limit", 10)
|
||||
|
||||
# Validate path
|
||||
allowed, error_msg = await self.check_path_allowed(path, tool_ctx)
|
||||
if not allowed:
|
||||
return error_msg
|
||||
|
||||
try:
|
||||
# Determine search scope
|
||||
project = self.project_manager.get_project_for_path(path)
|
||||
if not project:
|
||||
return "Error: No indexed project found for this path. Run 'vector --action index' first."
|
||||
|
||||
# Search
|
||||
await tool_ctx.info(f"Searching for: {query}")
|
||||
results = project.search(query, k=limit)
|
||||
|
||||
if not results:
|
||||
return f"No results found for: {query}"
|
||||
|
||||
# Format results
|
||||
output = [f"=== Vector Search Results for '{query}' ==="]
|
||||
output.append(f"Found {len(results)} matches\n")
|
||||
|
||||
for i, result in enumerate(results, 1):
|
||||
score = result.get("score", 0)
|
||||
file_path = result.get("file_path", "unknown")
|
||||
content = result.get("content", "")
|
||||
chunk_type = result.get("metadata", {}).get("type", "content")
|
||||
|
||||
output.append(f"Result {i} - Score: {score:.1%}")
|
||||
output.append(f"File: {file_path}")
|
||||
if chunk_type != "content":
|
||||
output.append(f"Type: {chunk_type}")
|
||||
output.append("-" * 60)
|
||||
|
||||
# Truncate content if too long
|
||||
if len(content) > 300:
|
||||
content = content[:300] + "..."
|
||||
output.append(content)
|
||||
output.append("")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Search failed: {str(e)}")
|
||||
return f"Error during search: {str(e)}"
|
||||
|
||||
async def _handle_index(self, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Handle indexing files."""
|
||||
path = params.get("path", ".")
|
||||
include = params.get("include")
|
||||
exclude = params.get("exclude")
|
||||
include_git = params.get("include_git", True)
|
||||
force = params.get("force_reindex", False)
|
||||
|
||||
# Validate path
|
||||
allowed, error_msg = await self.check_path_allowed(path, tool_ctx)
|
||||
if not allowed:
|
||||
return error_msg
|
||||
|
||||
try:
|
||||
await tool_ctx.info(f"Indexing {path}...")
|
||||
|
||||
# Get or create project
|
||||
project = self.project_manager.get_or_create_project(path)
|
||||
|
||||
# Index files
|
||||
stats = await project.index_directory(
|
||||
path,
|
||||
include_pattern=include,
|
||||
exclude_pattern=exclude,
|
||||
force_reindex=force,
|
||||
)
|
||||
|
||||
# Index git history if requested
|
||||
if include_git and Path(path).joinpath(".git").exists():
|
||||
await tool_ctx.info("Indexing git history...")
|
||||
git_stats = await project.index_git_history(path)
|
||||
stats["git_commits"] = git_stats.get("commits_indexed", 0)
|
||||
|
||||
# Format output
|
||||
output = [f"=== Indexing Complete ==="]
|
||||
output.append(f"Path: {path}")
|
||||
output.append(f"Files indexed: {stats.get('files_indexed', 0)}")
|
||||
output.append(f"Chunks created: {stats.get('chunks_created', 0)}")
|
||||
if stats.get("git_commits"):
|
||||
output.append(f"Git commits indexed: {stats['git_commits']}")
|
||||
output.append(
|
||||
f"Total documents: {project.get_stats().get('total_documents', 0)}"
|
||||
)
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Indexing failed: {str(e)}")
|
||||
return f"Error during indexing: {str(e)}"
|
||||
|
||||
async def _handle_stats(self, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Get vector store statistics."""
|
||||
path = params.get("path")
|
||||
|
||||
try:
|
||||
if path:
|
||||
# Stats for specific project
|
||||
project = self.project_manager.get_project_for_path(path)
|
||||
if not project:
|
||||
return f"No indexed project found for path: {path}"
|
||||
|
||||
stats = project.get_stats()
|
||||
output = [f"=== Vector Store Stats for {project.name} ==="]
|
||||
else:
|
||||
# Global stats
|
||||
stats = self.project_manager.get_global_stats()
|
||||
output = ["=== Global Vector Store Stats ==="]
|
||||
|
||||
output.append(f"Total documents: {stats.get('total_documents', 0)}")
|
||||
output.append(f"Total size: {stats.get('total_size_mb', 0):.1f} MB")
|
||||
|
||||
if stats.get("projects"):
|
||||
output.append(f"\nProjects indexed: {len(stats['projects'])}")
|
||||
for proj in stats["projects"]:
|
||||
output.append(
|
||||
f" - {proj['name']}: {proj['documents']} docs, {proj['size_mb']:.1f} MB"
|
||||
)
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to get stats: {str(e)}")
|
||||
return f"Error getting stats: {str(e)}"
|
||||
|
||||
async def _handle_clear(self, params: Dict[str, Any], tool_ctx) -> str:
|
||||
"""Clear vector store."""
|
||||
path = params.get("path")
|
||||
|
||||
if not path:
|
||||
return "Error: path is required for clear action"
|
||||
|
||||
# Validate path
|
||||
allowed, error_msg = await self.check_path_allowed(path, tool_ctx)
|
||||
if not allowed:
|
||||
return error_msg
|
||||
|
||||
try:
|
||||
project = self.project_manager.get_project_for_path(path)
|
||||
if not project:
|
||||
return f"No indexed project found for path: {path}"
|
||||
|
||||
# Get stats before clearing
|
||||
stats = project.get_stats()
|
||||
doc_count = stats.get("total_documents", 0)
|
||||
|
||||
# Clear
|
||||
project.clear()
|
||||
|
||||
return f"Cleared {doc_count} documents from vector store for {project.name}"
|
||||
|
||||
except Exception as e:
|
||||
await tool_ctx.error(f"Failed to clear: {str(e)}")
|
||||
return f"Error clearing vector store: {str(e)}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server."""
|
||||
pass
|
||||
@@ -1,119 +0,0 @@
|
||||
"""Vector indexing tool for adding documents to vector database."""
|
||||
|
||||
from typing import Dict, Unpack, Optional, TypedDict, final
|
||||
from pathlib import Path
|
||||
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, PermissionManager, auto_timeout
|
||||
|
||||
from .project_manager import ProjectVectorManager
|
||||
|
||||
|
||||
class VectorIndexParams(TypedDict, total=False):
|
||||
"""Parameters for vector indexing operations."""
|
||||
|
||||
file_path: str
|
||||
content: Optional[str]
|
||||
chunk_size: Optional[int]
|
||||
chunk_overlap: Optional[int]
|
||||
metadata: Optional[Dict[str, str]]
|
||||
|
||||
|
||||
@final
|
||||
class VectorIndexTool(BaseTool):
|
||||
"""Tool for indexing documents in the vector database."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
project_manager: ProjectVectorManager,
|
||||
):
|
||||
"""Initialize the vector index tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
project_manager: Project-aware vector store manager
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.project_manager = project_manager
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "vector_index"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Index documents in project-aware vector databases for semantic search.
|
||||
|
||||
Can index individual text content or entire files. Files are automatically assigned
|
||||
to the appropriate project database based on LLM.md detection or stored in the global
|
||||
database. Files are chunked for optimal search performance.
|
||||
|
||||
Projects are detected by finding LLM.md files, with databases stored in .hanzo/db
|
||||
directories alongside them. Use this to build searchable knowledge bases per project."""
|
||||
|
||||
@auto_timeout("vector_index")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[VectorIndexParams],
|
||||
) -> str:
|
||||
"""Index content or files in the vector database.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Indexing result message
|
||||
"""
|
||||
file_path = params.get("file_path")
|
||||
content = params.get("content")
|
||||
chunk_size = params.get("chunk_size", 1000)
|
||||
chunk_overlap = params.get("chunk_overlap", 200)
|
||||
metadata = params.get("metadata", {})
|
||||
|
||||
if not file_path and not content:
|
||||
return "Error: Either file_path or content must be provided"
|
||||
|
||||
try:
|
||||
if file_path:
|
||||
# Validate file access
|
||||
# Use permission manager's existing validation
|
||||
if not self.permission_manager.is_path_allowed(file_path):
|
||||
return f"Error: Access denied to path {file_path}"
|
||||
|
||||
if not Path(file_path).exists():
|
||||
return f"Error: File does not exist: {file_path}"
|
||||
|
||||
# Index file using project-aware manager
|
||||
doc_ids, project_info = (
|
||||
self.project_manager.add_file_to_appropriate_store(
|
||||
file_path=file_path,
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
metadata=metadata,
|
||||
)
|
||||
)
|
||||
|
||||
file_name = Path(file_path).name
|
||||
if project_info:
|
||||
return f"Successfully indexed {file_name} with {len(doc_ids)} chunks in project '{project_info.name}'"
|
||||
else:
|
||||
return f"Successfully indexed {file_name} with {len(doc_ids)} chunks in global database"
|
||||
|
||||
else:
|
||||
# Index content directly in global store (no project context)
|
||||
global_store = self.project_manager._get_global_store()
|
||||
doc_id = global_store.add_document(
|
||||
content=content,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
return f"Successfully indexed content as document {doc_id} in global database"
|
||||
|
||||
except Exception as e:
|
||||
return f"Error indexing content: {str(e)}"
|
||||
@@ -1,249 +0,0 @@
|
||||
"""Vector search tool for semantic document retrieval."""
|
||||
|
||||
import json
|
||||
from typing import List, Unpack, Optional, TypedDict, final
|
||||
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import BaseTool, PermissionManager, auto_timeout
|
||||
|
||||
from .project_manager import ProjectVectorManager
|
||||
|
||||
|
||||
class VectorSearchParams(TypedDict, total=False):
|
||||
"""Parameters for vector search operations."""
|
||||
|
||||
query: str
|
||||
limit: Optional[int]
|
||||
score_threshold: Optional[float]
|
||||
include_content: Optional[bool]
|
||||
file_filter: Optional[str]
|
||||
project_filter: Optional[List[str]]
|
||||
search_scope: Optional[str] # "all", "global", "current", or specific project name
|
||||
|
||||
|
||||
@final
|
||||
class VectorSearchTool(BaseTool):
|
||||
"""Tool for semantic search in the vector database."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
permission_manager: PermissionManager,
|
||||
project_manager: ProjectVectorManager,
|
||||
):
|
||||
"""Initialize the vector search tool.
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
project_manager: Project-aware vector store manager
|
||||
"""
|
||||
self.permission_manager = permission_manager
|
||||
self.project_manager = project_manager
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "vector_search"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Pure semantic/vector search using Infinity embedded database.
|
||||
|
||||
Searches indexed documents using vector embeddings to find semantically similar content.
|
||||
This is NOT keyword search - it finds documents based on meaning and context similarity.
|
||||
|
||||
Features:
|
||||
- Searches across project-specific vector databases
|
||||
- Returns similarity scores (0-1, higher is better)
|
||||
- Supports filtering by project or file
|
||||
- Automatically detects projects via LLM.md files
|
||||
|
||||
Use 'grep' for exact text/pattern matching, 'vector_search' for semantic similarity."""
|
||||
|
||||
@auto_timeout("vector_search")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
**params: Unpack[VectorSearchParams],
|
||||
) -> str:
|
||||
"""Search for similar documents in the vector database.
|
||||
|
||||
Args:
|
||||
ctx: MCP context
|
||||
**params: Tool parameters
|
||||
|
||||
Returns:
|
||||
Search results formatted as text
|
||||
"""
|
||||
query = params.get("query")
|
||||
if not query:
|
||||
return "Error: query parameter is required"
|
||||
|
||||
limit = params.get("limit", 10)
|
||||
score_threshold = params.get("score_threshold", 0.0)
|
||||
include_content = params.get("include_content", True)
|
||||
file_filter = params.get("file_filter")
|
||||
project_filter = params.get("project_filter")
|
||||
search_scope = params.get("search_scope", "all")
|
||||
|
||||
try:
|
||||
# Determine search strategy based on scope
|
||||
if search_scope == "all":
|
||||
# Search across all projects
|
||||
project_results = await self.project_manager.search_all_projects(
|
||||
query=query,
|
||||
limit_per_project=limit,
|
||||
score_threshold=score_threshold,
|
||||
include_global=True,
|
||||
project_filter=project_filter,
|
||||
)
|
||||
|
||||
# Combine and sort all results
|
||||
all_results = []
|
||||
for project_name, results in project_results.items():
|
||||
for result in results:
|
||||
# Add project info to metadata
|
||||
result.document.metadata = result.document.metadata or {}
|
||||
result.document.metadata["search_project"] = project_name
|
||||
all_results.append(result)
|
||||
|
||||
# Sort by score and limit
|
||||
all_results.sort(key=lambda x: x.score, reverse=True)
|
||||
results = all_results[:limit]
|
||||
|
||||
elif search_scope == "global":
|
||||
# Search only global store
|
||||
global_store = self.project_manager._get_global_store()
|
||||
results = global_store.search(
|
||||
query=query,
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
for result in results:
|
||||
result.document.metadata = result.document.metadata or {}
|
||||
result.document.metadata["search_project"] = "global"
|
||||
|
||||
else:
|
||||
# Search specific project or current context
|
||||
if search_scope != "current":
|
||||
# Search specific project by name
|
||||
project_info = None
|
||||
for _proj_key, proj_info in self.project_manager.projects.items():
|
||||
if proj_info.name == search_scope:
|
||||
project_info = proj_info
|
||||
break
|
||||
|
||||
if project_info:
|
||||
vector_store = self.project_manager.get_vector_store(
|
||||
project_info
|
||||
)
|
||||
results = vector_store.search(
|
||||
query=query,
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
for result in results:
|
||||
result.document.metadata = result.document.metadata or {}
|
||||
result.document.metadata["search_project"] = (
|
||||
project_info.name
|
||||
)
|
||||
else:
|
||||
return f"Project '{search_scope}' not found"
|
||||
else:
|
||||
# For "current", try to detect from working directory
|
||||
import os
|
||||
|
||||
current_dir = os.getcwd()
|
||||
project_info = self.project_manager.get_project_for_path(
|
||||
current_dir
|
||||
)
|
||||
|
||||
if project_info:
|
||||
vector_store = self.project_manager.get_vector_store(
|
||||
project_info
|
||||
)
|
||||
results = vector_store.search(
|
||||
query=query,
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
for result in results:
|
||||
result.document.metadata = result.document.metadata or {}
|
||||
result.document.metadata["search_project"] = (
|
||||
project_info.name
|
||||
)
|
||||
else:
|
||||
# Fall back to global store
|
||||
global_store = self.project_manager._get_global_store()
|
||||
results = global_store.search(
|
||||
query=query,
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
for result in results:
|
||||
result.document.metadata = result.document.metadata or {}
|
||||
result.document.metadata["search_project"] = "global"
|
||||
|
||||
if not results:
|
||||
return f"No results found for query: '{query}'"
|
||||
|
||||
# Filter by file if requested
|
||||
if file_filter:
|
||||
results = [
|
||||
r for r in results if file_filter in (r.document.file_path or "")
|
||||
]
|
||||
|
||||
# Format results
|
||||
output_lines = [f"Found {len(results)} results for query: '{query}'\n"]
|
||||
|
||||
for i, result in enumerate(results, 1):
|
||||
doc = result.document
|
||||
score_percent = result.score * 100
|
||||
|
||||
# Header with score and metadata
|
||||
project_name = doc.metadata.get("search_project", "unknown")
|
||||
header = f"Result {i} (Score: {score_percent:.1f}%) - Project: {project_name}"
|
||||
if doc.file_path:
|
||||
header += f" - {doc.file_path}"
|
||||
if doc.chunk_index is not None:
|
||||
header += f" [Chunk {doc.chunk_index}]"
|
||||
|
||||
output_lines.append(header)
|
||||
output_lines.append("-" * len(header))
|
||||
|
||||
# Add metadata if available
|
||||
if doc.metadata:
|
||||
relevant_metadata = {
|
||||
k: v
|
||||
for k, v in doc.metadata.items()
|
||||
if k not in ["chunk_number", "total_chunks", "search_project"]
|
||||
}
|
||||
if relevant_metadata:
|
||||
output_lines.append(
|
||||
f"Metadata: {json.dumps(relevant_metadata, indent=2)}"
|
||||
)
|
||||
|
||||
# Add content if requested
|
||||
if include_content:
|
||||
content = doc.content
|
||||
if len(content) > 500:
|
||||
content = content[:500] + "..."
|
||||
output_lines.append(f"Content:\n{content}")
|
||||
|
||||
output_lines.append("") # Empty line between results
|
||||
|
||||
return "\n".join(output_lines)
|
||||
|
||||
except Exception as e:
|
||||
return f"Error searching vector database: {str(e)}"
|
||||
|
||||
def register(self, mcp_server) -> None:
|
||||
"""Register this tool with the MCP server.
|
||||
|
||||
Args:
|
||||
mcp_server: The FastMCP server instance
|
||||
"""
|
||||
# This is a placeholder - the actual registration would happen
|
||||
# through the MCP server's tool registration mechanism
|
||||
pass
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hanzo-tools-vector"
|
||||
version = "0.2.1"
|
||||
version = "0.2.2"
|
||||
description = "Vector/embedding tools for Hanzo AI - indexing, search, RAG"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
@@ -19,10 +19,6 @@ dependencies = [
|
||||
"httpx>=0.25.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
full = ["sentence-transformers>=2.0.0", "faiss-cpu>=1.7.0"]
|
||||
infinity = ["infinity-embedded>=0.5.0"]
|
||||
|
||||
[project.entry-points."hanzo.tools"]
|
||||
vector = "hanzo_tools.vector:TOOLS"
|
||||
|
||||
|
||||
@@ -4,7 +4,10 @@ Import/registration checks plus offline tests of the cloud-backed VectorTool
|
||||
wiring (/v1/code/search, /v1/code/index, /v1/embeddings) and the mock-kill.
|
||||
"""
|
||||
|
||||
import re
|
||||
import asyncio
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -118,22 +121,35 @@ class TestCloudVectorTool:
|
||||
assert env["error"]["code"] == "INVALID_PARAMS"
|
||||
|
||||
|
||||
class TestNoSilentMock:
|
||||
def test_infinity_mock_is_opt_in(self):
|
||||
"""When infinity_embedded is absent, the store is UNAVAILABLE (no mock)."""
|
||||
from hanzo_tools.vector import infinity_store as ins
|
||||
class TestNoFake:
|
||||
"""The random-vector local store is deleted, not flag-guarded.
|
||||
|
||||
if ins.infinity_embedded is None:
|
||||
assert ins.INFINITY_AVAILABLE is False
|
||||
else:
|
||||
pytest.skip("infinity_embedded is installed in this environment")
|
||||
Regression guard: an importable local store means an embedder-less path
|
||||
exists again, and the only vectors it can produce are noise.
|
||||
"""
|
||||
|
||||
def test_local_embedding_raises_without_mock(self, monkeypatch):
|
||||
"""The local store must not silently return random vectors."""
|
||||
from hanzo_tools.vector.infinity_store import InfinityVectorStore
|
||||
@pytest.mark.parametrize(
|
||||
"gone",
|
||||
[
|
||||
"infinity_store",
|
||||
"mock_infinity",
|
||||
"vector_search",
|
||||
"vector_index",
|
||||
"index_tool",
|
||||
"git_ingester",
|
||||
"project_manager",
|
||||
],
|
||||
)
|
||||
def test_local_store_modules_are_gone(self, gone):
|
||||
with pytest.raises(ImportError):
|
||||
importlib.import_module(f"hanzo_tools.vector.{gone}")
|
||||
|
||||
monkeypatch.delenv("HANZO_VECTOR_ALLOW_MOCK", raising=False)
|
||||
store = InfinityVectorStore.__new__(InfinityVectorStore) # no __init__/DB
|
||||
store.dimension = 8
|
||||
with pytest.raises(NotImplementedError):
|
||||
store._generate_embedding("hello")
|
||||
def test_no_random_in_package(self):
|
||||
"""No module in the package may import `random` — vectors come from the service."""
|
||||
pkg = Path(importlib.import_module("hanzo_tools.vector").__file__).parent
|
||||
offenders = [
|
||||
p.name
|
||||
for p in pkg.glob("*.py")
|
||||
if re.search(r"^\s*(import random|from random import)", p.read_text(), re.M)
|
||||
]
|
||||
assert offenders == []
|
||||
|
||||
@@ -5,23 +5,61 @@ search, web search, vision — talks to it through this single seam. There is
|
||||
exactly one place that knows the base URL, the auth header, and how to turn a
|
||||
non-2xx into a typed error; tools compose it, never re-implement it.
|
||||
|
||||
Auth resolves in order: ``HANZO_API_KEY`` env, then the ``apiKey`` (hk- key) in
|
||||
``~/.hanzo/config.json``. Base URL is ``https://api.hanzo.ai``, overridable via
|
||||
``HANZO_API_BASE``.
|
||||
Auth resolves in order: ``HANZO_API_KEY`` env, the ``apiKey`` (hk- key) in
|
||||
``~/.hanzo/config.json``, then the ``hanzo`` CLI's live IAM session. Base URL is
|
||||
``https://api.hanzo.ai``, overridable via ``HANZO_API_BASE``.
|
||||
|
||||
Reference: HIP-0300 unified tools; api.hanzo.ai /v1 surface.
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import subprocess
|
||||
from typing import Any, ClassVar
|
||||
from pathlib import Path
|
||||
|
||||
DEFAULT_BASE = "https://api.hanzo.ai"
|
||||
|
||||
|
||||
# Memo for cli_session_token: shelling out costs ~100ms and the token is stable
|
||||
# for the session, so resolve it at most once per process. `False` means "asked
|
||||
# and there was none", which is distinct from "not asked yet" (None).
|
||||
_cli_token: str | None | bool = None
|
||||
|
||||
|
||||
def cli_session_token() -> str | None:
|
||||
"""The bearer token held by the ``hanzo`` CLI's IAM session.
|
||||
|
||||
The CLI already owns interactive login against hanzo.id, so asking it for a
|
||||
token means tools never prompt, never keep a second copy of the credential,
|
||||
and never parse a secret out of a file themselves.
|
||||
"""
|
||||
global _cli_token
|
||||
if _cli_token is not None:
|
||||
return _cli_token or None
|
||||
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["hanzo", "auth", "token"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=15,
|
||||
)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
_cli_token = False # no CLI on PATH: do not retry on every call
|
||||
return None
|
||||
|
||||
token = (out.stdout or "").strip()
|
||||
_cli_token = token if out.returncode == 0 and token else False
|
||||
return _cli_token or None
|
||||
|
||||
|
||||
def cloud_api_key() -> str | None:
|
||||
"""Resolve the hk- API key: env first, then ~/.hanzo/config.json."""
|
||||
"""Resolve the bearer credential.
|
||||
|
||||
Order: ``HANZO_API_KEY`` env, then the hk- key in ~/.hanzo/config.json,
|
||||
then the CLI's live IAM session.
|
||||
"""
|
||||
env = os.environ.get("HANZO_API_KEY") or os.environ.get("HANZO_KEY")
|
||||
if env and env.strip():
|
||||
return env.strip()
|
||||
@@ -33,8 +71,9 @@ def cloud_api_key() -> str | None:
|
||||
if key and str(key).strip():
|
||||
return str(key).strip()
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
return None
|
||||
pass
|
||||
|
||||
return cli_session_token()
|
||||
|
||||
|
||||
def cloud_api_base() -> str:
|
||||
@@ -136,6 +175,26 @@ class HanzoCloud:
|
||||
"""POST json_body to path, returning parsed JSON. Raises on failure."""
|
||||
return await self._request("POST", path, json=json_body or {})
|
||||
|
||||
async def call(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
json_body: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
"""Any-method call, for callers driving routes from the OpenAPI spec.
|
||||
|
||||
The spec names methods this client has no bespoke helper for (PUT,
|
||||
PATCH, DELETE), so it needs one generic seam rather than a helper per
|
||||
verb — the auth header and error mapping stay in this one place.
|
||||
"""
|
||||
kw: dict[str, Any] = {}
|
||||
if params:
|
||||
kw["params"] = {k: v for k, v in params.items() if v is not None}
|
||||
if json_body is not None:
|
||||
kw["json"] = json_body
|
||||
return await self._request(method.upper(), path, **kw)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
if self._client is not None:
|
||||
await self._client.aclose()
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hanzo-tools"
|
||||
version = "0.3.3"
|
||||
version = "0.3.4"
|
||||
description = "Hanzo AI tools - core infrastructure and tool bundles"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
|
||||
Reference in New Issue
Block a user