diff --git a/apis/collections/collections_helper.py b/apis/collections/collections_helper.py index 1971d53..dce7952 100644 --- a/apis/collections/collections_helper.py +++ b/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): diff --git a/apis/collections/collections_resource.py b/apis/collections/collections_resource.py index 4402e98..a17b9ca 100644 --- a/apis/collections/collections_resource.py +++ b/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/') 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: diff --git a/storages/chromadb_storage.py b/storages/chromadb_storage.py index 30ae453..8992e89 100644 --- a/storages/chromadb_storage.py +++ b/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)