From 05588767e88dacb4f1f51cabe728678e9d20ac10 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Tue, 30 Jan 2024 09:16:04 -0300
Subject: [PATCH] Refactor PineconeComponent in Pinecone.py

---
 .../components/vectorstores/Pinecone.py       | 46 ++++++++++++++-----
 1 file changed, 35 insertions(+), 11 deletions(-)

diff --git a/src/backend/langflow/components/vectorstores/Pinecone.py b/src/backend/langflow/components/vectorstores/Pinecone.py
index 1c04b6bc5..c8ac3c689 100644
--- a/src/backend/langflow/components/vectorstores/Pinecone.py
+++ b/src/backend/langflow/components/vectorstores/Pinecone.py
@@ -1,13 +1,12 @@
-from langflow import CustomComponent
-from typing import Optional, List, Union
-from langchain_community.vectorstores.pinecone import Pinecone
-from langflow.field_typing import (
-    Document,
-    Embeddings,
-)
+import os
+from typing import List, Optional, Union
+
+import pinecone  # type: ignore
 from langchain.schema import BaseRetriever
 from langchain.vectorstores.base import VectorStore
-import pinecone  # type: ignore
+from langchain_community.vectorstores.pinecone import Pinecone
+from langflow import CustomComponent
+from langflow.field_typing import Document, Embeddings
 
 
 class PineconeComponent(CustomComponent):
@@ -23,15 +22,40 @@ class PineconeComponent(CustomComponent):
             "pinecone_api_key": {"display_name": "Pinecone API Key", "default": "", "password": True, "required": True},
             "pinecone_env": {"display_name": "Pinecone Environment", "default": "", "required": True},
             "search_kwargs": {"display_name": "Search Kwargs", "default": "{}"},
+            "pool_threads": {"display_name": "Pool Threads", "default": 1, "advanced": True},
         }
 
     def build(
         self,
         embedding: Embeddings,
+        pinecone_env: str,
         documents: List[Document],
         index_name: Optional[str] = None,
         pinecone_api_key: Optional[str] = None,
-        pinecone_env: Optional[str] = None,
+        text_key: Optional[str] = "text",
+        namespace: Optional[str] = "default",
+        pool_threads: Optional[int] = None,
     ) -> Union[VectorStore, Pinecone, BaseRetriever]:
-        pinecone.init(api_key=pinecone_api_key, environment=pinecone_env)
-        return Pinecone.from_documents(documents=documents, embedding=embedding, index_name=index_name)
+        if pinecone_api_key is None or pinecone_env is None:
+            raise ValueError("Pinecone API Key and Environment are required.")
+        if os.getenv("PINECONE_API_KEY") is None and pinecone_api_key is None:
+            raise ValueError("Pinecone API Key is required.")
+
+        pinecone.init(api_key=pinecone_api_key, environment=pinecone_env)  # type: ignore
+        if documents:
+            return Pinecone.from_documents(
+                documents=documents,
+                embedding=embedding,
+                index_name=index_name,
+                pool_threads=pool_threads,
+                namespace=namespace,
+                text_key=text_key,
+            )
+
+        return Pinecone.from_existing_index(
+            index_name=index_name,
+            embedding=embedding,
+            text_key=text_key,
+            namespace=namespace,
+            pool_threads=pool_threads,
+        )