Advertisement
Ra7eN

AI chatbot with memory using SQLite3

Oct 3rd, 2024 (edited)
68
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 4.49 KB | Source Code | 0 0
  1. #
  2. # I have a site I am working on that has the whole thing including webpage tutorial (well you can cut/paste)
  3. # can find it here https://buildofflineaiwithmemory.neocities.org/
  4. # Updates: (contact rickscorpio64@gmail.com)
  5. #
  6. # Oct 7, 2024
  7. # ** WARNING ** REWROTE PYTHON SCRIPT BELOW.
  8. # Switched persistent memory to Annoy. Chromadb seemed to have issues
  9. # Now located here for full script and website ui integratoin
  10. # https://buildofflineaiwithmemory.neocities.org/
  11. #
  12. # Oct 3, 2024
  13. # system prompt is now external file (by request) called "personality.txt" for people that have huge peronality
  14. # date time added
  15. # FIX: "updating vector database: Collection conversations does not exist."
  16.  
  17. import time
  18. import ollama
  19. from fastapi import FastAPI
  20. from fastapi.middleware.cors import CORSMiddleware
  21. from fastapi.staticfiles import StaticFiles
  22. from fastapi.responses import PlainTextResponse, FileResponse
  23. from pydantic import BaseModel
  24. import uvicorn
  25. import os
  26. from annoy import AnnoyIndex
  27. import numpy as np
  28.  
  29. # Declare constants for the model names
  30. CHAT_MODEL = 'llama3'  # For generating text responses
  31. EMBEDDING_MODEL = 'nomic-embed-text'  # For generating 384-dimensional embeddings
  32. PERSONALITY_FILE = "personality.txt"  # The file containing the system prompt or personality
  33. ANNOY_INDEX_FILE = "memory.ann"
  34. VECTOR_DIM = 384  # Annoy index dimensionality
  35.  
  36. # Initialize FastAPI app
  37. app = FastAPI()
  38.  
  39. # CORS middleware setup
  40. app.add_middleware(
  41.     CORSMiddleware,
  42.     allow_origins=["*"],  # Allow all origins, or specify a domain.
  43.     allow_credentials=True,
  44.     allow_methods=["*"],  # Allow all HTTP methods.
  45.     allow_headers=["*"],  # Allow all headers.
  46. )
  47.  
  48. # Serve static files from the "static" directory
  49. app.mount("/static", StaticFiles(directory="static"), name="static")
  50.  
  51. # Initialize Annoy index
  52. annoy_index = AnnoyIndex(VECTOR_DIM, 'angular')
  53.  
  54. # Check if Annoy index exists and load it if present
  55. if os.path.exists(ANNOY_INDEX_FILE):
  56.     annoy_index.load(ANNOY_INDEX_FILE)
  57.  
  58. # Define the PromptModel
  59. class PromptModel(BaseModel):
  60.     prompt: str  # Ensure the prompt field is a string
  61.  
  62. # Load the system prompt/personality from the file
  63. def load_personality():
  64.     try:
  65.         with open(PERSONALITY_FILE, 'r') as file:
  66.             return file.read().strip()
  67.     except FileNotFoundError:
  68.         return "Default system instructions or personality prompt."
  69.  
  70. # Dummy function to simulate the chat model interaction (using llama3)
  71. def chat_with_model(system_prompt, user_prompt):
  72.     # This is where you'd integrate with the actual AI model, such as llama3.
  73.     # For now, it returns a dummy response.
  74.     return f"System says: {system_prompt}. You said: {user_prompt}"
  75.  
  76. # Store conversation and its embedding into Annoy index
  77. def store_conversation(prompt, response):
  78.     # Simulate generating an embedding from the prompt/response
  79.     embedding = np.random.rand(VECTOR_DIM).tolist()  # Dummy embedding generation
  80.  
  81.     # Add the new item to the in-memory Annoy index (rebuilding from scratch)
  82.     global annoy_index
  83.     annoy_index = AnnoyIndex(VECTOR_DIM, 'angular')  # Rebuild a fresh index
  84.     annoy_index.add_item(annoy_index.get_n_items(), embedding)
  85.    
  86.     # Save the updated index
  87.     annoy_index.build(10)  # Build the index with 10 trees
  88.     annoy_index.save(ANNOY_INDEX_FILE)
  89.  
  90. # Retrieve similar conversations (dummy implementation)
  91. def get_similar_conversations(prompt):
  92.     # Dummy similar conversations retrieval
  93.     return []
  94.  
  95. # FastAPI route to handle prompt input
  96. @app.post("/send_prompt/")
  97. def send_prompt(data: PromptModel):
  98.     prompt = data.prompt
  99.  
  100.     # Load personality/system prompt from file
  101.     system_prompt = load_personality()
  102.  
  103.     # Prepare the conversation with the system prompt and the user's input
  104.     convo = [
  105.         {'role': 'system', 'content': system_prompt},  # System instructions, only for AI's context
  106.         {'role': 'user', 'content': prompt}  # User input
  107.     ]
  108.  
  109.     # Get the AI response (using a proper model integration)
  110.     response = ollama.chat(CHAT_MODEL, convo)['message']['content']
  111.  
  112.     # Store the new conversation (user prompt and response)
  113.     store_conversation(prompt, response)
  114.  
  115.     # Return only the response, not the system prompt
  116.     return PlainTextResponse(response)
  117.  
  118. # Serve the index.html file at the root URL
  119. @app.get("/")
  120. def read_root():
  121.     return FileResponse('static/index.html')
  122.  
  123. # Start the FastAPI app
  124. if __name__ == "__main__":
  125.     uvicorn.run(app, host="127.0.0.1", port=8000)
  126.  
  127.  
Tags: ai AI memory
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement