refactor: move qdrant_client and embedding_model instantiation to indexing script

Author: vemonet

src/sparql_llm/agent/nodes/retrieval_docs.py

@@ -6,7 +6,8 @@
 
 from sparql_llm.agent.state import State, StepOutput
 from sparql_llm.agent.utils import get_msg_text
-from sparql_llm.config import Configuration, embedding_model, qdrant_client, settings
+from sparql_llm.config import Configuration, settings
+from sparql_llm.indexing.index_resources import embedding_model, qdrant_client
 
 # TODO: use grouping? https://qdrant.tech/documentation/concepts/search/#grouping-api
 # Which tools can I use for enrichment analysis?

src/sparql_llm/agent/nodes/retrieval_entities.py

@@ -7,7 +7,8 @@
 from qdrant_client import models
 
 from sparql_llm.agent.state import State, StepOutput
-from sparql_llm.config import Configuration, embedding_model, qdrant_client, settings
+from sparql_llm.config import Configuration, settings
+from sparql_llm.indexing.index_resources import embedding_model, qdrant_client
 
 # NOTE: experimental, not used in production
 

src/sparql_llm/agent/nodes/validation.py

@@ -10,7 +10,8 @@
 from sparql_llm.agent.prompts import FIX_QUERY_PROMPT
 from sparql_llm.agent.state import State, StepOutput
 from sparql_llm.config import Configuration, settings
-from sparql_llm.utils import endpoints_metadata, query_sparql
+from sparql_llm.indexing.index_resources import endpoints_metadata
+from sparql_llm.utils import query_sparql
 from sparql_llm.validate_sparql import validate_sparql_in_msg
 
 

src/sparql_llm/config.py

@@ -8,10 +8,8 @@
 from pathlib import Path
 from typing import Annotated, Any, Required, TypeVar
 
-from fastembed import TextEmbedding
 from langchain_core.runnables import RunnableConfig, ensure_config
 from pydantic_settings import BaseSettings, SettingsConfigDict
-from qdrant_client import QdrantClient
 from typing_extensions import TypedDict
 
 from sparql_llm.agent import prompts
@@ -224,20 +222,6 @@ def from_file(cls, filepath: str) -> Settings:
 settings = Settings.from_file(settings_filepath) if settings_filepath else Settings()
 # logger.info(f"📂 Using SETTINGS file: {settings_filepath}")
 
-# settings = Settings()
-
-# TODO: Getting `TypeError: cannot pickle '_thread.RLock' object` when doing `QdrantVectorStore.from_existing_collection(client=qdrant_client)`
-qdrant_client = (
-    QdrantClient(url=settings.vectordb_url, prefer_grpc=True, timeout=600)
-    if settings.vectordb_url.startswith(("http", "https"))
-    else QdrantClient(path=settings.vectordb_url)
-)
-
-embedding_model = TextEmbedding(
-    settings.embedding_model,
-    # providers=["CUDAExecutionProvider"], # Replace the fastembed dependency with fastembed-gpu to use your GPUs
-)
-
 
 # Configuration defined at runtime
 @dataclass(kw_only=True)

src/sparql_llm/indexing/index_entities.py

@@ -6,7 +6,8 @@
 from langchain_core.documents import Document
 from qdrant_client import models
 
-from sparql_llm.config import qdrant_client, settings
+from sparql_llm.config import settings
+from sparql_llm.indexing.index_resources import qdrant_client
 from sparql_llm.utils import query_sparql
 
 # NOTE: Run the script to extract entities from endpoints and generate embeddings for them (long):

src/sparql_llm/indexing/index_resources.py

@@ -3,20 +3,37 @@
 import httpx
 import pandas as pd
 from bs4 import BeautifulSoup
+from fastembed import TextEmbedding
 from langchain_core.documents import Document
 from markdownify import markdownify
-from qdrant_client import models
+from qdrant_client import QdrantClient, models
 from qdrant_client.http.models import Distance, VectorParams
 from rdflib import RDF, Dataset, Namespace
 
 from sparql_llm import SparqlExamplesLoader, SparqlInfoLoader, SparqlVoidShapesLoader
-from sparql_llm.config import SparqlEndpointLinks, embedding_model, qdrant_client, settings
+from sparql_llm.config import SparqlEndpointLinks, settings
 from sparql_llm.loaders.sparql_info_loader import GENERAL_INFO_DOC_TYPE
-from sparql_llm.utils import endpoints_metadata
+from sparql_llm.utils import EndpointsMetadataManager
 
 SCHEMA = Namespace("http://schema.org/")
 
 
+# Global instance, metadata loads lazily on first property access
+endpoints_metadata = EndpointsMetadataManager(settings.endpoints, settings.auto_init)
+
+# TODO: Getting `TypeError: cannot pickle '_thread.RLock' object` when doing `QdrantVectorStore.from_existing_collection(client=qdrant_client)`
+qdrant_client = (
+    QdrantClient(url=settings.vectordb_url, prefer_grpc=True, timeout=600)
+    if settings.vectordb_url.startswith(("http", "https"))
+    else QdrantClient(path=settings.vectordb_url)
+)
+
+embedding_model = TextEmbedding(
+    settings.embedding_model,
+    # providers=["CUDAExecutionProvider"], # Replace the fastembed dependency with fastembed-gpu to use your GPUs
+)
+
+
 def load_schemaorg_description(endpoint: SparqlEndpointLinks) -> list[Document]:
     """Extract datasets descriptions from the schema.org metadata in homepage of the endpoint"""
     docs = []

src/sparql_llm/mcp_server.py

@@ -4,12 +4,11 @@
 from mcp.server.fastmcp import FastMCP
 from qdrant_client.models import FieldCondition, Filter, MatchValue, ScoredPoint
 
-from sparql_llm.config import embedding_model, qdrant_client, settings
-from sparql_llm.indexing.index_resources import init_vectordb
-from sparql_llm.utils import endpoints_metadata, logger, query_sparql
+from sparql_llm.config import settings
+from sparql_llm.indexing.index_resources import embedding_model, endpoints_metadata, init_vectordb, qdrant_client
+from sparql_llm.utils import logger, query_sparql
 from sparql_llm.validate_sparql import validate_sparql
 
-logger
 # What are the rat orthologs of the human TP53?
 # TODO: MCP integrated https://github.com/modelcontextprotocol/python-sdk/pull/1007
 

src/sparql_llm/utils.py

@@ -7,7 +7,7 @@
 import httpx
 import rdflib
 
-from sparql_llm.config import SparqlEndpointLinks, settings
+from sparql_llm.config import SparqlEndpointLinks
 
 # Disable logger in your code with logging.getLogger("sparql_llm").setLevel(logging.WARNING)
 logger = logging.getLogger("sparql_llm")
@@ -255,7 +255,3 @@ def void_dict(self) -> "EndpointsSchemaDict":
     #     """Reset cached metadata (useful for re-initialization after init_vectordb)."""
     #     self._prefixes_map = {}
     #     self._void_dict = {}
-
-
-# Global instance, metadata loads lazily on first property access
-endpoints_metadata = EndpointsMetadataManager(settings.endpoints, settings.auto_init)

tests/benchmark_biodata.py

@@ -22,7 +22,8 @@
 
 # from sklearn.model_selection import KFold
 from sparql_llm import SparqlExamplesLoader, SparqlVoidShapesLoader
-from sparql_llm.config import embedding_model, qdrant_client, settings
+from sparql_llm.config import settings
+from sparql_llm.indexing.index_resources import embedding_model, qdrant_client
 from sparql_llm.utils import EndpointsMetadataManager, query_sparql
 from sparql_llm.validate_sparql import extract_sparql_queries