Skip to content

Commit 34f5b7b

Browse files
committed
expose MODEL_NAME
1 parent 71533f6 commit 34f5b7b

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

api/llm.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
KnowledgeGraphModelConfig
1313
)
1414

15-
#from graphrag_sdk.models.gemini import GeminiGenerativeModel
1615
from .prompts import (CYPHER_GEN_SYSTEM,
1716
CYPHER_GEN_PROMPT,
1817
GRAPH_QA_SYSTEM,
@@ -199,7 +198,9 @@ def _define_ontology() -> Ontology:
199198
def _create_kg_agent(repo_name: str):
200199
global ontology
201200

202-
model = LiteModel(model_name="gemini/gemini-2.0-flash-exp")
201+
model_name = os.getenv('MODEL_NAME', 'gemini/gemini-2.0-flash-exp')
202+
203+
model = LiteModel(model_name)
203204

204205
#ontology = _define_ontology()
205206
code_graph_kg = KnowledgeGraph(

0 commit comments

Comments
 (0)