add embeddings with openai

Author: VinciGit00

scrapegraphai/nodes/generate_answer_node_k_level.py

@@ -97,10 +97,24 @@ def execute(self, state: dict) -> dict:
 
         client = state["vectorial_db"]
 
-        answer_db = client.query(
-            collection_name="vectorial_collection",
-            query_text=state["question"]
+        if state.get("embeddings"):
+            import openai
+            openai_client = openai.Client()
+
+            answer_db = client.search(
+            collection_name="collection",
+            query_vector=openai_client.embeddings.create(
+                input=["What is the best to use for vector search scaling?"],
+                model=state.get("embeddings").get("model"),
+            )
+            .data[0]
+            .embedding,
         )
+        else:
+            answer_db = client.query(
+                collection_name="vectorial_collection",
+                query_text=state["question"]
+            )
 
         ## TODO: from the id get the data
         results_db = [elem for elem in state[answer_db]]

scrapegraphai/nodes/rag_node.py

@@ -4,6 +4,7 @@
 from typing import List, Optional
 from .base_node import BaseNode
 from qdrant_client import QdrantClient
+from qdrant_client.models import PointStruct, VectorParams, Distance
 
 class RAGNode(BaseNode):
     """
@@ -52,6 +53,41 @@ def execute(self, state: dict) -> dict:
         docs = [elem.get("summary") for elem in state.get("descriptions", {})]
         ids = [elem.get("id") for elem in state.get("descriptions", {})]
 
+        if state.get("embeddings"):
+            import openai
+            openai_client = openai.Client()
+
+            files = state.get("documents")
+
+            array_of_embeddings = []
+            i=0
+
+            for file in files:
+                embeddings = openai_client.embeddings.create(input=file,
+                                                             model=state.get("embeddings").get("model"))
+                i+=1
+                points = PointStruct(
+                        id=i,
+                        vector=embeddings,
+                        payload={"text": file},
+                    )
+
+                array_of_embeddings.append(points)
+
+            collection_name = "collection"
+
+            client.create_collection(
+                collection_name,
+                vectors_config=VectorParams(
+                    size=1536,
+                    distance=Distance.COSINE,
+                ),
+            )
+            client.upsert(collection_name, points)
+
+            state["vectorial_db"] = client
+            return state
+
         client.add(
             collection_name="vectorial_collection",
             documents=docs,