Refactor PineconeComponent in Pinecone.py
This commit is contained in:
parent
910c55fa0c
commit
05588767e8
1 changed files with 35 additions and 11 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue