Skip to content

Adapters

magetools.adapters

Adapters for Magetools.

This module provides concrete implementations of the embedding and vector store protocols. Dependencies (chromadb, google-genai) are lazily imported to allow the core magetools package to be used with alternative providers.

ChromaVectorStore

Bases: VectorStoreProtocol

Adapter for ChromaDB.

Source code in src/magetools/adapters.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
class ChromaVectorStore(VectorStoreProtocol):
    """Adapter for ChromaDB."""

    def __init__(self, path: str):
        chromadb, _ = _import_chromadb()
        self.client = chromadb.PersistentClient(path=str(path))

    def get_collection(self, name: str, embedding_function: Any) -> Any:
        return self.client.get_collection(
            name=name, embedding_function=embedding_function
        )

    def list_collections(self) -> list[Any]:
        return self.client.list_collections()

    def get_or_create_collection(self, name: str, embedding_function: Any) -> Any:
        return self.client.get_or_create_collection(
            name=name, embedding_function=embedding_function
        )

    async def close(self) -> None:
        """Close database connection."""
        self.client = None

GoogleGenAIProvider

Bases: EmbeddingProviderProtocol

Provider for Google Generative AI embeddings.

Source code in src/magetools/adapters.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
class GoogleGenAIProvider(EmbeddingProviderProtocol):
    """Provider for Google Generative AI embeddings."""

    def __init__(self, config: MageToolsConfig | None = None):
        self.config = config or get_config()
        genai = _import_genai()
        self.client = genai.Client()
        self._genai = genai

    def get_embedding_function(self) -> Any:
        _, embedding_functions = _import_chromadb()
        return embedding_functions.GoogleGenerativeAiEmbeddingFunction(
            model_name=self.config.embedding_model, task_type="SEMANTIC_SIMILARITY"
        )

    def generate_content(self, prompt: str) -> str:
        """Generates content using Google Gemini model."""
        logger.debug(f"Generating content with model {self.config.model_name}...")
        response = self.client.models.generate_content(
            model=self.config.model_name, contents=prompt
        )
        logger.debug("Generated content.")
        return response.text

    async def close(self) -> None:
        """Cleanup AI client."""
        pass

MockEmbeddingProvider

Bases: EmbeddingProviderProtocol

Mock provider for when google-genai is not available.

This allows users to try the library without credentials. Search functionality will be degraded but the app won't crash.

Source code in src/magetools/adapters.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
class MockEmbeddingProvider(EmbeddingProviderProtocol):
    """Mock provider for when google-genai is not available.

    This allows users to try the library without credentials.
    Search functionality will be degraded but the app won't crash.
    """

    def __init__(self, config: MageToolsConfig | None = None):
        self.config = config or get_config()
        logger.warning(
            "Using MockEmbeddingProvider - search functionality will be limited. "
            "Install google-genai and set GOOGLE_API_KEY for full functionality."
        )

    def get_embedding_function(self) -> Any:
        """Return a simple mock embedding function."""
        return _MockEmbeddingFunction()

    def generate_content(self, prompt: str) -> str:
        """Return a placeholder summary."""
        logger.debug("MockProvider: Returning placeholder summary")
        return "Summary not available (MockProvider - no LLM configured)"

    async def close(self) -> None:
        """No-op cleanup."""
        pass