fix: replace eval() with json.loads(), add embedding retry, use dynamic version (#23)

Signed-off-by: Cheney Zhang <chen.zhang@zilliz.com>

Author: zc277584121

src/vector_graph_rag/api/app.py

@@ -20,7 +20,7 @@
 from fastapi.responses import FileResponse
 from pydantic import BaseModel, Field
 
-from vector_graph_rag import VectorGraphRAG
+from vector_graph_rag import VectorGraphRAG, __version__
 from vector_graph_rag.config import Settings, get_settings
 from vector_graph_rag.storage.milvus import MilvusStore
 from vector_graph_rag.graph.graph import Graph
@@ -32,7 +32,7 @@
 class HealthResponse(BaseModel):
     """Health check response."""
     status: str = Field(default="ok", description="Service status")
-    version: str = Field(default="0.1.0", description="API version")
+    version: str = Field(default=__version__, description="API version")
 
 
 class GraphInfo(BaseModel):
@@ -258,7 +258,7 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
     app = FastAPI(
         title="Vector Graph RAG API",
         description="Graph RAG using pure vector search with Milvus",
-        version="0.1.0",
+        version=__version__,
     )
 
     # Add CORS middleware for frontend
@@ -300,7 +300,7 @@ def get_graph(graph_name: Optional[str] = None) -> Graph:
     @app.get("/health", response_model=HealthResponse, tags=["System"])
     async def health_check():
         """Check if the service is running."""
-        return HealthResponse(status="ok", version="0.1.0")
+        return HealthResponse(status="ok", version=__version__)
 
     @app.get("/graphs", response_model=ListGraphsResponse, tags=["System"])
     async def list_graphs():

src/vector_graph_rag/llm/extractor.py

@@ -304,10 +304,10 @@ def _load_tsv_cache(self, cache_file: str) -> None:
                 query = row.get(query_col, '')
                 triples_str = row.get('triples', '{}')
                 try:
-                    triples_data = eval(triples_str) if isinstance(triples_str, str) else triples_str
+                    triples_data = json.loads(triples_str) if isinstance(triples_str, str) else triples_str
                     if isinstance(triples_data, dict) and 'named_entities' in triples_data:
                         self.ner_tsv_cache[query] = triples_data['named_entities']
-                except:
+                except (json.JSONDecodeError, KeyError, TypeError):
                     pass
             print(f"Loaded {len(self.ner_tsv_cache)} NER entries from {cache_file}")
         except Exception as e:

src/vector_graph_rag/storage/embeddings.py

@@ -156,10 +156,26 @@ def __init__(self, model_name: str, api_key: str, base_url: Optional[str] = None
 
         self.model_name = model_name
         self.client = OpenAI(api_key=api_key, base_url=base_url)
-        self._retry_decorator = retry(
+
+        # Wrap the API call with retry logic
+        @retry(
             stop=stop_after_attempt(3),
             wait=wait_exponential(multiplier=1, min=2, max=10),
         )
+        def _call_api(texts):
+            return self.client.embeddings.create(model=self.model_name, input=texts)
+
+        self._call_api = _call_api
+
+        # Detect embedding dimension lazily
+        self._dimension: Optional[int] = None
+
+    def _get_dimension(self) -> int:
+        """Get embedding dimension by making a test call."""
+        if self._dimension is None:
+            response = self._call_api(["test"])
+            self._dimension = len(response.data[0].embedding)
+        return self._dimension
 
     def encode(
         self,
@@ -182,9 +198,9 @@ def encode(
         valid_texts = [texts[i] for i in valid_indices]
 
         if not valid_texts:
-            return np.zeros((len(texts), 1536))
+            return np.zeros((len(texts), self._get_dimension()))
 
-        response = self.client.embeddings.create(model=self.model_name, input=valid_texts)
+        response = self._call_api(valid_texts)
         sorted_data = sorted(response.data, key=lambda x: x.index)
         valid_embeddings = np.array([item.embedding for item in sorted_data])