From 33908da18edddaaeec145b2d2031073ce4cf931c Mon Sep 17 00:00:00 2001
From: Pavlo Paliychuk <paul@getzep.com>
Date: Thu, 19 Sep 2024 16:35:36 -0400
Subject: [PATCH] feat: add OpenAI configuration options to Settings and update
 LLM client setup (#126)

* chore: Add model and base url customization support for graphiti svc

* fix: formatter
---
 server/graph_service/config.py       | 3 +++
 server/graph_service/zep_graphiti.py | 6 ++++++
 2 files changed, 9 insertions(+)

diff --git a/server/graph_service/config.py b/server/graph_service/config.py
index e80d0586..f3082a5f 100644
--- a/server/graph_service/config.py
+++ b/server/graph_service/config.py
@@ -2,11 +2,14 @@ from functools import lru_cache
 from typing import Annotated
 
 from fastapi import Depends
+from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict  # type: ignore
 
 
 class Settings(BaseSettings):
     openai_api_key: str
+    openai_base_url: str | None = Field(None)
+    model_name: str | None = Field(None)
     neo4j_uri: str
     neo4j_user: str
     neo4j_password: str
diff --git a/server/graph_service/zep_graphiti.py b/server/graph_service/zep_graphiti.py
index 236a1b35..c1ff566f 100644
--- a/server/graph_service/zep_graphiti.py
+++ b/server/graph_service/zep_graphiti.py
@@ -54,6 +54,12 @@ async def get_graphiti(settings: ZepEnvDep):
         user=settings.neo4j_user,
         password=settings.neo4j_password,
     )
+    if settings.openai_base_url is not None:
+        client.llm_client.config.base_url = settings.openai_base_url
+    if settings.openai_api_key is not None:
+        client.llm_client.config.api_key = settings.openai_api_key
+    if settings.model_name is not None:
+        client.llm_client.model = settings.model_name
     try:
         yield client
     finally: