Update Pinecone.py to use the new Pinecone package and add support for BaseRetriever

This commit is contained in:
anovazzi1 2024-01-16 16:11:12 -03:00
commit cba57bb8f2

View file

@ -1,13 +1,13 @@
from langflow import CustomComponent from langflow import CustomComponent
from typing import Optional, List from typing import Optional, List, Union
from langchain.vectorstores import Pinecone from langchain_community.vectorstores.pinecone import Pinecone
from langflow.field_typing import ( from langflow.field_typing import (
Document, Document,
Embeddings, Embeddings,
NestedDict,
) )
from langchain.schema import BaseRetriever
from langchain.vectorstores.base import VectorStore
import pinecone
class PineconeComponent(CustomComponent): class PineconeComponent(CustomComponent):
display_name = "Pinecone" display_name = "Pinecone"
description = "Construct Pinecone wrapper from raw documents." description = "Construct Pinecone wrapper from raw documents."
@ -28,17 +28,8 @@ class PineconeComponent(CustomComponent):
embedding: Embeddings, embedding: Embeddings,
documents: Optional[List[Document]] = None, documents: Optional[List[Document]] = None,
index_name: Optional[str] = None, index_name: Optional[str] = None,
namespace: Optional[str] = None,
pinecone_api_key: Optional[str] = None, pinecone_api_key: Optional[str] = None,
pinecone_env: Optional[str] = None, pinecone_env: Optional[str] = None,
search_kwargs: Optional[NestedDict] = None, ) -> Union[VectorStore,Pinecone,BaseRetriever]:
) -> Pinecone: pinecone.init(api_key=pinecone_api_key,environment=pinecone_env)
return Pinecone( return Pinecone.from_documents(documents=documents,embedding=embedding,index_name=index_name)
documents=documents,
embedding=embedding,
index_name=index_name,
namespace=namespace,
pinecone_api_key=pinecone_api_key,
pinecone_env=pinecone_env,
search_kwargs=search_kwargs,
)