diff --git a/pyproject.toml b/pyproject.toml index 9bcd8d5..5765fa1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,3 +54,7 @@ select = ["E", "F", "I", "UP", "B"] python_version = "3.11" strict = true warn_unused_configs = true + +[[tool.mypy.overrides]] +module = ["chromadb", "chromadb.*"] +ignore_missing_imports = true diff --git a/src/tai/chroma_telemetry.py b/src/tai/chroma_telemetry.py index 310c65e..0e81c3c 100644 --- a/src/tai/chroma_telemetry.py +++ b/src/tai/chroma_telemetry.py @@ -9,7 +9,6 @@ from __future__ import annotations from chromadb.config import System from chromadb.telemetry.product import ProductTelemetryClient, ProductTelemetryEvent -from overrides import override class NoOpProductTelemetryClient(ProductTelemetryClient): @@ -18,7 +17,6 @@ class NoOpProductTelemetryClient(ProductTelemetryClient): def __init__(self, system: System): super().__init__(system) - @override def capture(self, event: ProductTelemetryEvent) -> None: del event return None \ No newline at end of file diff --git a/src/tai/cli.py b/src/tai/cli.py index 8dbd50e..22b048a 100644 --- a/src/tai/cli.py +++ b/src/tai/cli.py @@ -411,7 +411,7 @@ async def _interactive_loop( else: line = sys.stdin.readline() # non-TTY / piped mode if not line: - return + return last_response command = line.strip() console.print(f"\n[bold cyan]tai[/bold cyan][dim] >[/dim] {command}") except (EOFError, KeyboardInterrupt): diff --git a/src/tai/runbook_store.py b/src/tai/runbook_store.py index 42778e6..4f34969 100644 --- a/src/tai/runbook_store.py +++ b/src/tai/runbook_store.py @@ -17,7 +17,7 @@ from __future__ import annotations import re from dataclasses import dataclass, field from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from tai.ai_client import AIClient @@ -123,7 +123,7 @@ class RunbookStore: self._client = chromadb.PersistentClient(path=str(path)) else: self._client = chromadb.PersistentClient(path=str(path), settings=settings) - self._collection = self._client.get_or_create_collection( + self._collection: Any = self._client.get_or_create_collection( name=_COLLECTION_NAME, metadata={"hnsw:space": "cosine"}, ) @@ -241,11 +241,14 @@ class RunbookStore: return [] results = self._collection.get(include=["metadatas"]) metas = results.get("metadatas") or [] - return [dict(m) for m in metas] + entries: list[dict[str, str]] = [] + for meta in metas: + entries.append({str(k): str(v) for k, v in dict(meta).items()}) + return entries def count(self) -> int: """Return the number of indexed runbook documents.""" - return self._collection.count() + return int(self._collection.count()) # --------------------------------------------------------------------------- diff --git a/src/tai/session_store.py b/src/tai/session_store.py index a2d3564..63ef633 100644 --- a/src/tai/session_store.py +++ b/src/tai/session_store.py @@ -5,7 +5,7 @@ from __future__ import annotations from dataclasses import dataclass from datetime import UTC, datetime from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from tai.ai_client import AIClient @@ -48,14 +48,14 @@ class SessionStore: self._client = chromadb.PersistentClient(path=str(path)) else: self._client = chromadb.PersistentClient(path=str(path), settings=settings) - self._collection = self._client.get_or_create_collection( + self._collection: Any = self._client.get_or_create_collection( name=_COLLECTION_NAME, metadata={"hnsw:space": "cosine"}, ) def count(self) -> int: """Return number of indexed session summaries.""" - return self._collection.count() + return int(self._collection.count()) def index_session(self, host: str, issue: str, summary: str, ai: AIClient) -> str: """Embed and upsert one session summary into persistent storage."""