From 8e8c267ff30659867421781aa83a853f19425e09 Mon Sep 17 00:00:00 2001 From: fegloff Date: Fri, 22 Dec 2023 11:59:00 -0400 Subject: [PATCH 1/2] add document endpoint --- apis/collections/collections_helper.py | 11 ++++++ apis/collections/collections_resource.py | 48 +++++++++++++++++++++--- storages/chromadb_storage.py | 11 ++++-- 3 files changed, 61 insertions(+), 9 deletions(-) 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..61dd222 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') +@api.route('/document/add') 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. + 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) From 9e874b6ae02933369131ad5a6fd178f495ba9e70 Mon Sep 17 00:00:00 2001 From: fegloff Date: Tue, 26 Dec 2023 14:04:33 -0400 Subject: [PATCH 2/2] update document/add endpoint class --- apis/collections/collections_resource.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apis/collections/collections_resource.py b/apis/collections/collections_resource.py index 61dd222..a17b9ca 100644 --- a/apis/collections/collections_resource.py +++ b/apis/collections/collections_resource.py @@ -189,7 +189,7 @@ def delete(self, collection_name): return make_response(jsonify({"error": "An unexpected error occurred."}), 500) @api.route('/document/add') -class WebCrawlerTextRes(Resource): +class DocumentAddRes(Resource): # # @copy_current_request_context def post(self):