upgrade graphrag-sdk and move to litellm

.env.template

@@ -3,6 +3,6 @@ FALKORDB_PORT=<port>
 FALKORDB_USERNAME=<username>
 FALKORDB_PASSWORD=<password>
 OPENAI_API_KEY=<openai_api_key>
-GOOGLE_API_KEY=<google_api_key>
+GEMINI_API_KEY=<gemini_api_key>
 SECRET_TOKEN=<secret_token>
 

README.md

@@ -3,19 +3,23 @@
 [![Discord](https://img.shields.io/discord/1146782921294884966?style=flat-square)](https://discord.com/invite/6M4QwDXn2w)
 
 ## Getting Started
+
 [Live Demo](https://code-graph.falkordb.com/)
 
 ## Running locally
 
 ### Run FalkorDB
+
 Free cloud instance: https://app.falkordb.cloud/signup
 
 Or by running locally with docker:
+
 ```bash
 docker run -p 6379:6379 -p 3000:3000 -it --rm falkordb/falkordb:latest
 ```
 
 ### Config
+
 Create your own `.env` file from the `.env.template` file
 
 Start the server:
@@ -24,16 +28,21 @@ flask --app api/index.py run --debug
 ```
 
 ### Creating a graph
+
 Process a local source folder:
+
 ```bash
 curl -X POST http://127.0.0.1:5000/analyze_folder -H "Content-Type: application/json" -d '{"path": "<FULL_PATH_TO_FOLDER>", "ignore": [<OPTIONAL_IGNORE_LIST>]}' -H "Authorization: <.ENV_SECRET_TOKEN>"
 ```
 
 For example:
+
 ```bash
 curl -X POST http://127.0.0.1:5000/analyze_folder -H "Content-Type: application/json" -d '{"path": "/Users/roilipman/Dev/GraphRAG-SDK", "ignore": ["./.github", "./build"]}' -H "Authorization: OpenSesame"
 ```
+
 ## Working with your graph
+
 Once the source code analysis completes your FalkorDB DB will be populated with
 a graph representation of your source code, the graph name should be the same as
 the name of the folder you've requested to analyze, for the example above a graph named:
@@ -43,6 +52,8 @@ At the moment only the Python and C languages are supported, we do intend to sup
 
 At this point you can explore and query your source code using various tools
 Here are several options:
-1. FalkorDB built-in UI
+
+1. [Code-Graph UI](https://github.com/FalkorDB/code-graph)
+1. FalkorDB [Browser](https://github.com/FalkorDB/falkordb-browser/)
 2. One of FalkorDB's [clients](https://docs.falkordb.com/clients.html)
 3. Use FalkorDB [GraphRAG-SDK](https://github.com/FalkorDB/GraphRAG-SDK) to connect an LLM for natural language exploration.

api/graph.py

@@ -28,7 +28,7 @@ def get_repos() -> List[str]:
                   password=os.getenv('FALKORDB_PASSWORD', None))
 
     graphs = db.list_graphs()
-    graphs = [g for g in graphs if not g.endswith('_git')]
+    graphs = [g for g in graphs if not (g.endswith('_git') or g.endswith('_schema'))]
     return graphs
 
 class Graph():

api/llm.py

@@ -1,7 +1,7 @@
 import os
 import logging
 
-from graphrag_sdk.models.openai import OpenAiGenerativeModel
+from graphrag_sdk.models.litellm import LiteModel
 from graphrag_sdk import (
     Ontology,
     Entity,
@@ -12,7 +12,6 @@
     KnowledgeGraphModelConfig
 )
 
-#from graphrag_sdk.models.gemini import GeminiGenerativeModel
 from .prompts import (CYPHER_GEN_SYSTEM,
                      CYPHER_GEN_PROMPT,
                      GRAPH_QA_SYSTEM,
@@ -199,15 +198,15 @@ def _define_ontology() -> Ontology:
 def _create_kg_agent(repo_name: str):
     global ontology
 
-    openapi_model    = OpenAiGenerativeModel("gpt-4o")
-    #gemini_model     = GeminiGenerativeModel("gemini-1.5-flash-001")
-    #gemini_model_pro = GeminiGenerativeModel("gemini-1.5-pro")
+    model_name = os.getenv('MODEL_NAME', 'gemini/gemini-2.0-flash-exp')
+
+    model = LiteModel(model_name)
 
     #ontology = _define_ontology()
     code_graph_kg = KnowledgeGraph(
         name=repo_name,
         ontology=ontology,
-        model_config=KnowledgeGraphModelConfig.with_model(openapi_model),
+        model_config=KnowledgeGraphModelConfig.with_model(model),
         host=os.getenv('FALKORDB_HOST', 'localhost'),
         port=os.getenv('FALKORDB_PORT', 6379),
         username=os.getenv('FALKORDB_USERNAME', None),