Update app.py

app.py

@@ -5,6 +5,11 @@ import torch
 import optimum
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from fastapi.responses import HTMLResponse
+from datetime import datetime
+import random
+import string
+from datasets import Dataset
+import json
 
 # Set environment variables for GPU usage and memory allocation
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
@@ -14,6 +19,7 @@ torch.cuda.set_per_process_memory_fraction(0.8)  # Adjust the fraction as needed
 # Initialize FastAPI application
 app = FastAPI(root_path="/api/v1")
 
+
 # Load the model and tokenizer
 model_name_or_path = "TheBloke/Wizard-Vicuna-7B-Uncensored-GPTQ"
 
@@ -74,6 +80,51 @@ def hermes_model():
 
 model, tokenizer = hermes_model()
 
+def generate_id(length=5):
+    """
+    Generates a random alphanumeric ID.
+    
+    Args:
+        length (int): The length of the ID.
+    
+    Returns:
+        str: A random alphanumeric ID.
+    """
+    return ''.join(random.choices(string.ascii_letters + string.digits, k=length))
+
+def generate_thread_id():
+    """
+    Generates a unique thread ID for each conversation.
+    
+    Returns:
+        str: A unique thread ID.
+    """
+    return generate_id()
+
+def generate_message_id():
+    """
+    Generates a random alphanumeric message ID.
+    
+    Returns:
+        str: A random alphanumeric message ID.
+    """
+    return generate_id()
+
+
+def save_conversation(user_id, conversation):
+    hf_space_path="articko/conversations"
+    """
+    Save conversation history to disk.
+    
+    Args:
+        user_id (str): The unique identifier for the user.
+        conversation (dict): The conversation data.
+        hf_space_path (str): The path to the Hugging Face Space.
+    """
+    with open(f'{hf_space_path}/conversations.jsonl', 'a') as file:
+        json.dump({user_id: conversation}, file)
+        file.write('\n')
+
 def chat_response(msg_prompt: str) -> dict:
     """
     Generates a response from the model given a prompt.
@@ -202,6 +253,46 @@ async def hermes_chat(request: Request):
         raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
+
+@app.post('/chat_thread/{user_id}')
+async def chat_thread(request: Request, user_id: str):
+    """
+    Starts a new conversation thread with a provided prompt for a specific user.
+
+    Args:
+        request (Request): The HTTP request object containing the user prompt.
+        user_id (str): The unique identifier for the user.
+
+    Returns:
+        dict: The response generated by the model along with the user's conversation history.
+    """
+    try:
+        thread_id = generate_thread_id()
+
+        data = await request.json()
+        msg_prompt = data.get('msg_prompt')
+
+        if not msg_prompt:
+            raise HTTPException(status_code=400, detail="Prompt not provided")
+
+        # Generate response
+        response = chat_response(msg_prompt)
+
+        # Generate message ID
+        message_id = generate_message_id()
+
+        # Construct conversation entry
+        conversation_entry = {'thread_id': thread_id, 'message_id': message_id, 'user': msg_prompt, 'assistant': response}
+
+        # Save conversation history to disk
+        save_conversation(user_id, conversation_entry)
+
+        # Return response and thread ID
+        return {'response': conversation_entry}
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
         
 @app.post('/prompted_chat')
 async def prompted_chat(request: Request):