DH-4917/fix the embedding issue

dataherald/sql_generator/dataherald_sqlagent.py

@@ -22,6 +22,7 @@
     CallbackManagerForToolRun,
 )
 from langchain.chains.llm import LLMChain
+from langchain.embeddings import OpenAIEmbeddings
 from langchain.schema import AgentAction
 from langchain.tools.base import BaseTool
 from overrides import override
@@ -205,9 +206,10 @@ def get_embedding(
         self, text: str, model: str = "text-embedding-ada-002"
     ) -> List[float]:
         text = text.replace("\n", " ")
-        return openai.Embedding.create(input=[text], model=model)["data"][0][
-            "embedding"
-        ]
+        embedding = OpenAIEmbeddings(
+            openai_api_key=os.environ.get("OPENAI_API_KEY"), model=model
+        )
+        return embedding.embed_query(text)
 
     def cosine_similarity(self, a: List[float], b: List[float]) -> float:
         return round(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)), 4)

requirements.txt

@@ -34,3 +34,4 @@ sphinx-book-theme==1.0.1
 boto3==1.28.38
 botocore==1.31.38
 PyAthena==3.0.6
+tiktoken==0.5.1