open router fix

nimh-dsst · Aug 29, 2024 · 486d5ea · 486d5ea
1 parent 1d10af8
commit 486d5ea
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/external_components/llm_extraction/app.py b/external_components/llm_extraction/app.py
@@ -10,13 +10,12 @@
 
 # from pydantic import BaseModel, Field
 from osm.schemas.metrics_schemas import LLMExtractorMetrics
-
+from llama_index.llms.openrouter import OpenRouter
 LLM_MODELS = {
-    "gpt-4o-2024-08-06": OpenAI(
-        base_url="https://openrouter.ai/api/v1",
-        api_key=os.environ.get("OPENROUTER_API_KEY"),
-        model="gpt-4o-2024-08-06",
-    )
+    "gpt-4o-2024-08-06": OpenRouter(
+    api_key=os.environ["OPENROUTER_API_KEY"],
+    model="openai/chatgpt-4o-latest",
+)
 }