diff --git a/main.py b/main.py index 1d25171..447f31c 100644 --- a/main.py +++ b/main.py @@ -1,12 +1,23 @@ +import os +from pyexpat import model +import sys from fastapi import FastAPI from pydantic import BaseModel import subprocess +import json +from pymongo import MongoClient app = FastAPI() class InputData(BaseModel): + user_id: str query: str +# MONGO_URI = os.getenv("MONGO_URI", "mongodb://localhost:27017/") +client = MongoClient("mongodb+srv://trmnteam:tBp54siAioeGkVpb@cluster0.68spx5t.mongodb.net/?retryWrites=true&w=majority&appName=Cluster0") +db = client["chat_db"] +history_collection = db["chat_history"] + @app.get("/") def read_root(): return {"message": "Hello, world!"} @@ -30,13 +41,56 @@ def run_script(): @app.post("/ask-ai") def ask_ai(data: InputData): try: - result = subprocess.run(['python', 'chat.py', data.query], capture_output=True, text=True) - - return {"response": result.stdout.strip(), "error": result.stderr.strip()} - + result = subprocess.run( + ['python', 'chat.py', data.user_id, data.query], + capture_output=True, + text=True + ) + response_text = result.stdout.strip() + error_text = result.stderr.strip() + + # Load existing history + history = load_history(data.user_id) + + # Append the user's input and the AI's response to the history + history.append({"role": "user", "parts": [data.query]}) + history.append({"role": "model", "parts": [response_text]}) + + # Save the updated history + save_history(data.user_id, history) + + return {"response": response_text, "error": error_text} except Exception as e: return {"error": str(e)} if __name__ == "__main__": import uvicorn - uvicorn.run(app, host="0.0.0.0", port=8000) \ No newline at end of file + uvicorn.run(app, host="0.0.0.0", port=8000) + +# db connection and other configurations + +def load_history(user_id): + doc = history_collection.find_one({"user_id": user_id}) + if doc and "history" in doc: + return doc["history"] + return [] + +def save_history(user_id, history): + history_collection.update_one( + {"user_id": user_id}, + {"$set": {"history": history}}, + upsert=True + ) + +if len(sys.argv) > 2: + user_id = sys.argv[1] + user_input = sys.argv[2] + history = load_history(user_id) + chat_session = model.start_chat(history=history) + response = chat_session.send_message(user_input) + history.append({"role": "user", "parts": [user_input]}) + history.append({"role": "model", "parts": [response.text]}) + save_history(user_id, history) + print(response.text) +else: + print("Error: No user_id or input provided.") \ No newline at end of file