Add index name parameter to agent (#68)

hive_agent/agent.py

@@ -64,6 +64,7 @@ def __init__(
         retrieve=False,
         required_exts=supported_exts,
         retrieval_tool="basic",
+        index_name : Optional[str] = None,
         load_index_file=False,
         swarm_mode=False,
         chat_only_mode=False,
@@ -91,6 +92,7 @@ def __init__(
         self.retrieve = retrieve
         self.required_exts = required_exts
         self.retrieval_tool = retrieval_tool
+        self.index_name = index_name
         self.load_index_file = load_index_file
         logging.basicConfig(stream=sys.stdout, level=self.__config.get("log"))
         logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
@@ -277,17 +279,26 @@ def add_batch_indexes(self):
 
         if "chroma" in self.retrieval_tool:
             chroma_retriever = ChromaRetriever()
-            index, file_names = chroma_retriever.create_index()
+            if self.index_name is not None:
+                index, file_names = chroma_retriever.create_index(collection_name=self.index_name)
+            else:
+                index, file_names = chroma_retriever.create_index()
             self.index_store.add_index(chroma_retriever.name, index, file_names)
 
         if "pinecone-serverless" in self.retrieval_tool:
             pinecone_retriever = PineconeRetriever()
-            index, file_names = pinecone_retriever.create_serverless_index()
+            if self.index_name is not None:
+                index, file_names = pinecone_retriever.create_serverless_index(collection_name=self.index_name)
+            else:
+                index, file_names = pinecone_retriever.create_serverless_index()
             self.index_store.add_index(pinecone_retriever.name, index, file_names)
 
         if "pinecone-pod" in self.retrieval_tool:
             pinecone_retriever = PineconeRetriever()
-            index, file_names = pinecone_retriever.create_pod_index()
+            if self.index_name is not None:
+                index, file_names = pinecone_retriever.create_pod_index(collection_name=self.index_name)
+            else:
+                index, file_names = pinecone_retriever.create_pod_index()
             self.index_store.add_index(pinecone_retriever.name, index, file_names)
 
             self.index_store.save_to_file()

hive_agent/tools/retriever/pinecone_retrieve.py

@@ -34,7 +34,7 @@ def create_serverless_index(
         folder_path=None,
         prefix='',
         bucket=None,
-        name="hive-agent-pinecone",
+        collection_name="hive-agent-pinecone",
         dimension=1536,
         metric="euclidean",
         cloud="aws",
@@ -45,12 +45,12 @@ def create_serverless_index(
         else:
             documents, file_names = self._load_documents(file_path, folder_path)
         self.pinecone_client.create_index(
-            name=name,
+            name=collection_name,
             dimension=dimension,
             metric=metric,
             spec=ServerlessSpec(cloud=cloud, region=region),
         )
-        pinecone_index = self.pinecone_client.Index(name)
+        pinecone_index = self.pinecone_client.Index(collection_name)
 
         vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
         storage_context = StorageContext.from_defaults(vector_store=vector_store)
@@ -65,7 +65,7 @@ def create_pod_index(
         folder_path=None,
         s3_endpoint_url=None,
         bucket=None,
-        name="hive-agent-pinecone-pod",
+        collection_name="hive-agent-pinecone-pod",
         dimension=1536,
         metric="cosine",
         environment="us-east1-gcp",
@@ -77,12 +77,12 @@ def create_pod_index(
         else:
             documents, file_names = self._load_documents(file_path, folder_path)
         self.pinecone_client.create_index(
-            name=name,
+            name=collection_name,
             dimension=dimension,
             metric=metric,
             spec=PodSpec(environment=environment, pod_type=pod_type, pods=pods),
         )
-        pinecone_index = self.pinecone_client.Index(name)
+        pinecone_index = self.pinecone_client.Index(collection_name)
 
         vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
         storage_context = StorageContext.from_defaults(vector_store=vector_store)