Refactor PineconeComponent in Pinecone.py

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-01-30 09:16:04 -03:00
commit 05588767e8

View file

@ -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,
)