From dacca690033d0197d878bee7fbbb6adb0b0373b1 Mon Sep 17 00:00:00 2001
From: rajandangi <rajandangi649@gmail.com>
Date: Thu, 2 May 2024 22:53:27 +0930
Subject: [PATCH] feat: add support for customizable embedding model in OpenAI
 integration

This commit introduces the ability to specify an embedding model when using OpenAI's embedding API. It updates the configuration files and the embedding component to accept a user-defined embedding model. This enhancement provides flexibility in selecting different models for various use cases.
---
 private_gpt/components/embedding/embedding_component.py | 5 ++++-
 private_gpt/settings/settings.py                        | 4 ++++
 settings-openai.yaml                                    | 1 +
 settings.yaml                                           | 1 +
 4 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/private_gpt/components/embedding/embedding_component.py b/private_gpt/components/embedding/embedding_component.py
index 77e8c3d..c2d67d2 100644
--- a/private_gpt/components/embedding/embedding_component.py
+++ b/private_gpt/components/embedding/embedding_component.py
@@ -56,7 +56,10 @@ class EmbeddingComponent:
                     ) from e
 
                 openai_settings = settings.openai.api_key
-                self.embedding_model = OpenAIEmbedding(api_key=openai_settings)
+                self.embedding_model = OpenAIEmbedding(
+                    api_key=openai_settings,
+                    model=openai_settings.embedding_model,
+                )
             case "ollama":
                 try:
                     from llama_index.embeddings.ollama import (  # type: ignore
diff --git a/private_gpt/settings/settings.py b/private_gpt/settings/settings.py
index c4c5e20..e2ebc8d 100644
--- a/private_gpt/settings/settings.py
+++ b/private_gpt/settings/settings.py
@@ -201,6 +201,10 @@ class OpenAISettings(BaseModel):
         description="Base URL of OpenAI API. Example: 'https://api.openai.com/v1'.",
     )
     api_key: str
+    embedding_model: str = Field(
+        "text-embedding-ada-002",
+        description="OpenAI Model to use. Example: 'text-embedding-ada-002'.",
+    )
     model: str = Field(
         "gpt-3.5-turbo",
         description="OpenAI Model to use. Example: 'gpt-4'.",
diff --git a/settings-openai.yaml b/settings-openai.yaml
index 093fcea..d929b39 100644
--- a/settings-openai.yaml
+++ b/settings-openai.yaml
@@ -3,6 +3,7 @@ server:
 
 llm:
   mode: openai
+  embedding_model: text-embedding-ada-002
 
 embedding:
   mode: openai
diff --git a/settings.yaml b/settings.yaml
index d8d2500..52b4238 100644
--- a/settings.yaml
+++ b/settings.yaml
@@ -95,6 +95,7 @@ sagemaker:
 openai:
   api_key: ${OPENAI_API_KEY:}
   model: gpt-3.5-turbo
+  embedding_model: text-embedding-ada-002
 
 ollama:
   llm_model: llama2