Skip to content

Commit 20488ba

Browse files
committed
upgrade graphrag-sdk and move to litellm
1 parent c20267a commit 20488ba

File tree

3 files changed

+1443
-382
lines changed

3 files changed

+1443
-382
lines changed

api/llm.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -199,15 +199,13 @@ def _define_ontology() -> Ontology:
199199
def _create_kg_agent(repo_name: str):
200200
global ontology
201201

202-
openapi_model = OpenAiGenerativeModel("gpt-4o")
203-
#gemini_model = GeminiGenerativeModel("gemini-1.5-flash-001")
204-
#gemini_model_pro = GeminiGenerativeModel("gemini-1.5-pro")
202+
model = LiteModel(model_name="gemini/gemini-2.0-flash-exp")
205203

206204
#ontology = _define_ontology()
207205
code_graph_kg = KnowledgeGraph(
208206
name=repo_name,
209207
ontology=ontology,
210-
model_config=KnowledgeGraphModelConfig.with_model(openapi_model),
208+
model_config=KnowledgeGraphModelConfig.with_model(model),
211209
host=os.getenv('FALKORDB_HOST', 'localhost'),
212210
port=os.getenv('FALKORDB_PORT', 6379),
213211
username=os.getenv('FALKORDB_USERNAME', None),

0 commit comments

Comments
 (0)