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 77e8c3d4..c2d67d28 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 c4c5e20d..e2ebc8d5 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 093fcea9..d929b398 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 d8d2500c..52b4238a 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