From 3f58a17e47605bac18ad72cfe84a01efba1fdce3 Mon Sep 17 00:00:00 2001 From: PVBLIC Foundation Date: Sat, 10 May 2025 06:07:27 -0700 Subject: [PATCH] Update pinecone.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit • Removed the unused Pinecone REST‐client import; we now only import ServerlessSpec and the gRPC client. • Enhanced close() • Call self.client.close() to explicitly shut down the underlying gRPC channel. • Log success or a warning on failure. • Still tear down the thread‐pool executor afterward. • Context‐manager support • Added __enter__()/__exit__() so you can do: with PineconeClient() as client: client.insert(...) # automatically calls client.close() --- .../open_webui/retrieval/vector/dbs/pinecone.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/backend/open_webui/retrieval/vector/dbs/pinecone.py b/backend/open_webui/retrieval/vector/dbs/pinecone.py index 22aa11cb4..2f44e70e7 100644 --- a/backend/open_webui/retrieval/vector/dbs/pinecone.py +++ b/backend/open_webui/retrieval/vector/dbs/pinecone.py @@ -1,7 +1,7 @@ from typing import Optional, List, Dict, Any, Union import logging import time # for measuring elapsed time -from pinecone import Pinecone, ServerlessSpec +from pinecone import ServerlessSpec import asyncio # for async upserts import functools # for partial binding in async tasks @@ -496,5 +496,18 @@ class PineconeClient(VectorDBBase): raise def close(self): - """Shut down the thread pool.""" + """Shut down the gRPC channel and thread pool.""" + try: + self.client.close() + log.info("Pinecone gRPC channel closed.") + except Exception as e: + log.warning(f"Failed to close Pinecone gRPC channel: {e}") self._executor.shutdown(wait=True) + + def __enter__(self): + """Enter context manager.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Exit context manager, ensuring resources are cleaned up.""" + self.close()