add document endpoint

apis/collections/collections_helper.py

@@ -34,6 +34,17 @@ def collection_query(self, collection_name, prompt, conversation):
         else: 
             raise InvalidCollectionName("Collection Error", "collection doesn't exist, please try again later", 404)
 
+    def collection_add_document(self, collection_name, text):
+        collection = self.get_collection(collection_name)
+        if collection:
+            doc_id = self.db.generate_document_id(collection_name)
+            collection.add(ids=[doc_id], embeddings=[text])
+            return {
+                "price": 0.02
+            }
+        else:
+            raise InvalidCollectionName("Collection Error", "collection doesn't exist, please try again later", 404)
+
     def delete_collection(self, collection_name):
         collection = self.db.get_existing_collection(collection_name)
         if (collection):

apis/collections/collections_resource.py

@@ -87,6 +87,44 @@ def __collection_request_handler(self, url, collection_name, file_name, context)
             error = CollectionError(dict( collection_name = collection_name))
             error.save()
 
+@api.route('/query')
+class WebCrawlerTextRes(Resource):
+    # 
+    # @copy_current_request_context
+    def post(self):
+        """
+        Endpoint to handle LLMs request.
+        Receives a message from the user, processes it, and returns a response from the model.
+        """ 
+        data = request.json
+        prompt = data.get('prompt')
+        collection_name = data.get('collectionName')
+        url = data.get('url')
+        conversation = data.get('conversation')
+        chat_history = [ChatMessage(content=item.get('content'), role=item.get('role')) for item in conversation]
+        try:
+            current_app.logger.info(f'Inquiring a collection {collection_name}')
+            if collection_name:
+                response = collection_helper.collection_query(collection_name, prompt, chat_history) 
+                return make_response(jsonify(response), 200)
+            else:
+                current_app.logger.error('Bad request')
+                return make_response(jsonify({"error": "Bad request"}), 400)
+        except InvalidCollectionName as e:
+
+            current_app.logger.error(e)
+            return make_response(jsonify({"error": e.args[1]}), 404)   
+        except OpenAIError as e:
+            # Handle OpenAI API errors
+            error_message = str(e)
+            current_app.logger.error(f"OpenAI API Error: {error_message}")
+            return jsonify({"error": error_message}), 500
+        except Exception as e:
+            error_message = str(e)
+            current_app.logger.error(f'ERROR ***************: {error_message}')
+            current_app.logger.error(e)
+            return make_response(jsonify({"error": "An unexpected error occurred."}), 500)
+
 @api.route('/document/<collection_name>')
 class CheckDocument(Resource):
 
@@ -150,14 +188,13 @@ def delete(self, collection_name):
             current_app.logger.error(f"Unexpected Error: {error_message}")
             return make_response(jsonify({"error": "An unexpected error occurred."}), 500)
 
-@api.route('/query')
-class WebCrawlerTextRes(Resource):
+@api.route('/document/add')
+class DocumentAddRes(Resource):
     # 
     # @copy_current_request_context
     def post(self):
         """
-        Endpoint to handle LLMs request.
-        Receives a message from the user, processes it, and returns a response from the model.
+        Endpoint to handle to add a document to a collection.
         """ 
         data = request.json
         prompt = data.get('prompt')
@@ -168,13 +205,12 @@ def post(self):
         try:
             current_app.logger.info(f'Inquiring a collection {collection_name}')
             if collection_name:
-                response = collection_helper.collection_query(collection_name, prompt, chat_history) 
+                response = collection_helper.collection_add_document(collection_name=collection_name, text=chat_history) 
                 return make_response(jsonify(response), 200)
             else:
                 current_app.logger.error('Bad request')
                 return make_response(jsonify({"error": "Bad request"}), 400)
-        except InvalidCollectionName as e:
-
+        except InvalidCollectionName as e:    
             current_app.logger.error(e)
             return make_response(jsonify({"error": e.args[1]}), 404)   
         except OpenAIError as e:

storages/chromadb_storage.py

@@ -2,10 +2,12 @@
 from llama_index import Document, VectorStoreIndex, LLMPredictor, ServiceContext, PromptHelper
 from llama_index.vector_stores import ChromaVectorStore
 from llama_index.storage.storage_context import StorageContext
-import chromadb
+from langchain.chat_models import ChatOpenAI
 from chromadb.config import Settings
 import hashlib
-from langchain.chat_models import ChatOpenAI
+import datetime
+import chromadb
+
 from config import config
 class ChromaStorage:
 
@@ -34,7 +36,10 @@ def get_collection_name(self, chat_id, url):
         valid_characters = ''.join(c for c in hashed if c.isalnum() or c in ('_', '-'))
         return f"chat{chat_id}-{valid_characters}"
 
-
+    def generate_document_id(self, collection_name):
+        now = datetime.datetime.now()
+        now_str = now.strftime("%Y%m%d%H%M%S")
+        return f"{collection_name}-{now_str}"
     def get_existing_collection(self, collection_name):
         try:
             collection = self.db.get_collection(collection_name)