144 lines
6.1 KiB
Python
144 lines
6.1 KiB
Python
from typing import Any, Dict, List
|
|
|
|
from pinecone import PineconeAsyncio
|
|
|
|
from letta.constants import PINECONE_CLOUD, PINECONE_EMBEDDING_MODEL, PINECONE_METRIC, PINECONE_REGION, PINECONE_TEXT_FIELD_NAME
|
|
from letta.log import get_logger
|
|
from letta.schemas.user import User
|
|
from letta.settings import settings
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def should_use_pinecone(verbose: bool = False):
|
|
if verbose:
|
|
logger.info(
|
|
"Pinecone check: enable_pinecone=%s, api_key=%s, agent_index=%s, source_index=%s",
|
|
settings.enable_pinecone,
|
|
bool(settings.pinecone_api_key),
|
|
bool(settings.pinecone_agent_index),
|
|
bool(settings.pinecone_source_index),
|
|
)
|
|
|
|
return settings.enable_pinecone and settings.pinecone_api_key and settings.pinecone_agent_index and settings.pinecone_source_index
|
|
|
|
|
|
async def upsert_pinecone_indices():
|
|
from pinecone import IndexEmbed, PineconeAsyncio
|
|
|
|
for index_name in get_pinecone_indices():
|
|
async with PineconeAsyncio(api_key=settings.pinecone_api_key) as pc:
|
|
if not await pc.has_index(index_name):
|
|
await pc.create_index_for_model(
|
|
name=index_name,
|
|
cloud=PINECONE_CLOUD,
|
|
region=PINECONE_REGION,
|
|
embed=IndexEmbed(model=PINECONE_EMBEDDING_MODEL, field_map={"text": PINECONE_TEXT_FIELD_NAME}, metric=PINECONE_METRIC),
|
|
)
|
|
|
|
|
|
def get_pinecone_indices() -> List[str]:
|
|
return [settings.pinecone_agent_index, settings.pinecone_source_index]
|
|
|
|
|
|
async def upsert_file_records_to_pinecone_index(file_id: str, source_id: str, chunks: List[str], actor: User):
|
|
records = []
|
|
for i, chunk in enumerate(chunks):
|
|
record = {
|
|
"_id": f"{file_id}_{i}",
|
|
PINECONE_TEXT_FIELD_NAME: chunk,
|
|
"file_id": file_id,
|
|
"source_id": source_id,
|
|
}
|
|
records.append(record)
|
|
|
|
return await upsert_records_to_pinecone_index(records, actor)
|
|
|
|
|
|
async def delete_file_records_from_pinecone_index(file_id: str, actor: User):
|
|
from pinecone.exceptions.exceptions import NotFoundException
|
|
|
|
namespace = actor.organization_id
|
|
try:
|
|
async with PineconeAsyncio(api_key=settings.pinecone_api_key) as pc:
|
|
description = await pc.describe_index(name=settings.pinecone_source_index)
|
|
async with pc.IndexAsyncio(host=description.index.host) as dense_index:
|
|
await dense_index.delete(
|
|
filter={
|
|
"file_id": {"$eq": file_id},
|
|
},
|
|
namespace=namespace,
|
|
)
|
|
except NotFoundException:
|
|
logger.warning(f"Pinecone namespace {namespace} not found for {file_id} and {actor.organization_id}")
|
|
|
|
|
|
async def delete_source_records_from_pinecone_index(source_id: str, actor: User):
|
|
from pinecone.exceptions.exceptions import NotFoundException
|
|
|
|
namespace = actor.organization_id
|
|
try:
|
|
async with PineconeAsyncio(api_key=settings.pinecone_api_key) as pc:
|
|
description = await pc.describe_index(name=settings.pinecone_source_index)
|
|
async with pc.IndexAsyncio(host=description.index.host) as dense_index:
|
|
await dense_index.delete(filter={"source_id": {"$eq": source_id}}, namespace=namespace)
|
|
except NotFoundException:
|
|
logger.warning(f"Pinecone namespace {namespace} not found for {source_id} and {actor.organization_id}")
|
|
|
|
|
|
async def upsert_records_to_pinecone_index(records: List[dict], actor: User):
|
|
async with PineconeAsyncio(api_key=settings.pinecone_api_key) as pc:
|
|
description = await pc.describe_index(name=settings.pinecone_source_index)
|
|
async with pc.IndexAsyncio(host=description.index.host) as dense_index:
|
|
await dense_index.upsert_records(actor.organization_id, records)
|
|
|
|
|
|
async def search_pinecone_index(query: str, limit: int, filter: Dict[str, Any], actor: User) -> Dict[str, Any]:
|
|
async with PineconeAsyncio(api_key=settings.pinecone_api_key) as pc:
|
|
description = await pc.describe_index(name=settings.pinecone_source_index)
|
|
async with pc.IndexAsyncio(host=description.index.host) as dense_index:
|
|
namespace = actor.organization_id
|
|
|
|
try:
|
|
# Search the dense index with reranking
|
|
search_results = await dense_index.search(
|
|
namespace=namespace,
|
|
query={
|
|
"top_k": limit,
|
|
"inputs": {"text": query},
|
|
"filter": filter,
|
|
},
|
|
rerank={"model": "bge-reranker-v2-m3", "top_n": limit, "rank_fields": [PINECONE_TEXT_FIELD_NAME]},
|
|
)
|
|
return search_results
|
|
except Exception as e:
|
|
logger.warning(f"Failed to search Pinecone namespace {actor.organization_id}: {str(e)}")
|
|
raise e
|
|
|
|
|
|
async def list_pinecone_index_for_files(file_id: str, actor: User, limit: int = None, pagination_token: str = None) -> List[str]:
|
|
from pinecone.exceptions.exceptions import NotFoundException
|
|
|
|
namespace = actor.organization_id
|
|
try:
|
|
async with PineconeAsyncio(api_key=settings.pinecone_api_key) as pc:
|
|
description = await pc.describe_index(name=settings.pinecone_source_index)
|
|
async with pc.IndexAsyncio(host=description.index.host) as dense_index:
|
|
|
|
kwargs = {"namespace": namespace, "prefix": file_id}
|
|
if limit is not None:
|
|
kwargs["limit"] = limit
|
|
if pagination_token is not None:
|
|
kwargs["pagination_token"] = pagination_token
|
|
|
|
try:
|
|
result = []
|
|
async for ids in dense_index.list(**kwargs):
|
|
result.extend(ids)
|
|
return result
|
|
except Exception as e:
|
|
logger.warning(f"Failed to list Pinecone namespace {actor.organization_id}: {str(e)}")
|
|
raise e
|
|
except NotFoundException:
|
|
logger.warning(f"Pinecone namespace {namespace} not found for {file_id} and {actor.organization_id}")
|