From 31e8c9628b9bae5ef24582339867be8c5fa6b8ab Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Thu, 7 Dec 2023 22:36:56 -0800 Subject: [PATCH] Update mediator and cache and add routes --- .gitignore | 2 + cache.py | 24 ++++++++--- main.py | 91 +++++++++++++++++++++++++++++++++++------ mediator.py | 115 +++++++++++++++++++++++++++++++++++++++------------- 4 files changed, 185 insertions(+), 47 deletions(-) diff --git a/.gitignore b/.gitignore index 68bc17f9..dbce267a 100644 --- a/.gitignore +++ b/.gitignore @@ -158,3 +158,5 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + +.DS_Store diff --git a/cache.py b/cache.py index fb48f24e..dc399ac5 100644 --- a/cache.py +++ b/cache.py @@ -5,7 +5,7 @@ in OrderedDict data structure. from collections import OrderedDict from mediator import SupabaseMediator import uuid -from typing import List +from typing import List, Dict from langchain.schema import BaseMessage, Document from pydantic import BaseModel import sentry_sdk @@ -13,24 +13,36 @@ import sentry_sdk class Conversation: "Wrapper Class for storing contexts between channels. Using an object to pass by reference avoid additional cache hits" @sentry_sdk.trace - def __init__(self, mediator: SupabaseMediator, user_id: str, conversation_id: str = str(uuid.uuid4()), location_id: str = "web"): + def __init__(self, mediator: SupabaseMediator, user_id: str, session_id: str = str(uuid.uuid4()), location_id: str = "web", metadata: Dict = {}): self.mediator: SupabaseMediator = mediator self.user_id: str = user_id - self.conversation_id: str = conversation_id + self.session_id: str = session_id self.location_id: str = location_id + self.metadata: Dict = metadata @sentry_sdk.trace def add_message(self, message_type: str, message: BaseMessage,) -> None: - self.mediator.add_message(self.conversation_id, self.user_id, message_type, message) + self.mediator.add_message(self.session_id, self.user_id, message_type, message) @sentry_sdk.trace def messages(self, message_type: str) -> List[BaseMessage]: - return self.mediator.messages(self.conversation_id, self.user_id, message_type) + return self.mediator.messages(self.session_id, self.user_id, message_type) + + @sentry_sdk.trace + def delete(self) -> None: + self.mediator.delete_session(self.session_id) + + @sentry_sdk.trace + def restart(self) -> None: + self.delete() + representation = self.mediator.add_session(user_id=self.user_id, location_id=self.location_id) + self.session_id: str = representation["id"] + self.metadata = representation["metadata"] # vector DB fn @sentry_sdk.trace def add_texts(self, texts: List[str]) -> None: - metadatas = [{"conversation_id": self.conversation_id, "user_id": self.user_id} for _ in range(len(texts))] + metadatas = [{"session_id": self.session_id, "user_id": self.user_id} for _ in range(len(texts))] self.mediator.vector_table.add_texts(texts, metadatas) # vector DB fn diff --git a/main.py b/main.py index f55348d2..9f6e247e 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,7 @@ from fastapi import FastAPI, HTTPException -from fastapi.responses import StreamingResponse +from fastapi.responses import StreamingResponse, JSONResponse from pydantic import BaseModel +from typing import Optional, Dict import asyncio # Local @@ -25,31 +26,97 @@ app = FastAPI() MEDIATOR = SupabaseMediator() LOCK = asyncio.Lock() -class ConversationInput(BaseModel): +class SessionInput(BaseModel): user_id: str - conversation_id: str + session_id: str message: str + message_type: str @app.get("/") def root(): return {"message": "Hello World"} -@app.post("/chat") -async def voe(inp: ConversationInput): +class Session(BaseModel): + session_id: str + user_id: str + location_id: Optional[str] + metadata: Optional[Dict] + +### Session Interface Routes ### +@app.get('/messages') +async def get_messages(inp: Session): async with LOCK: - conversation = Conversation(MEDIATOR, user_id=inp.user_id, conversation_id=inp.conversation_id) - if conversation is None: + return MEDIATOR.messages(inp.session_id, inp.user_id, 'response') + +@app.post('/messages/add') +async def add_message(inp: SessionInput): + async with LOCK: + MEDIATOR.add_message(inp.session_id, inp.user_id, inp.message_type, inp.message) + +### Session Meta Routes ### +class UnknownSession(BaseModel): + user_id: str + location_id: Optional[str] + +@app.get('/session') +async def get_session(inp: UnknownSession): + """Return session ids and metadata associated with a user and location""" + location = "default" if inp.location_id is None else inp.location_id + async with LOCK: + id: str = MEDIATOR.session(location, inp.user_id) + return JSONResponse(status_code=200, content={"session_id": id}) + +@app.post('/session/add') +async def add_session(inp: UnknownSession): + location = "default" if inp.location_id is None else inp.location_id + async with LOCK: + MEDIATOR.add_session(location, inp.user_id) + return JSONResponse(status_code=200, content={"message": "OK"}) + +class SessionMeta(BaseModel): + session_id: str + user_id: str + metadata: Optional[Dict] + +@app.delete('/session/delete') +async def delete_session(inp: SessionMeta): + """Delete a specific session""" + async with LOCK: + MEDIATOR.delete_session(inp.user_id, inp.session_id) + return JSONResponse(status_code=200, content={"message": "OK"}) + +@app.patch('/session/update') +async def update_session(inp: SessionMeta): + async with LOCK: + MEDIATOR.update_session(inp.user_id, inp.session_id, inp.metadata) + return JSONResponse(status_code=200, content={"message": "OK"}) + +### API Fundamentals +# @app.post('/tenant/new') +# @app.post('application/new') +# @app.get('/tenant') + + +## Honcho Utilities +# @app.get('/theoryofmind') +# @app.get('/voe') + +@app.post("/chat") +async def voe(inp: SessionInput): + async with LOCK: + session = Conversation(MEDIATOR, user_id=inp.user_id, session_id=inp.session_id) + if session is None: raise HTTPException(status_code=404, detail="Item not found") - thought, response = await BloomChain.chat(conversation, inp.message) + thought, response = await BloomChain.chat(session, inp.message) return { "thought": thought, "response": response } @app.post("/stream") -async def stream(inp: ConversationInput): +async def stream(inp: SessionInput): async with LOCK: - conversation = Conversation(MEDIATOR, user_id=inp.user_id, conversation_id=inp.conversation_id) - if conversation is None: + session = Conversation(MEDIATOR, user_id=inp.user_id, session_id=inp.session_id) + if session is None: raise HTTPException(status_code=404, detail="Item not found") - return StreamingResponse(BloomChain.stream(conversation, inp.message)) + return StreamingResponse(BloomChain.stream(session, inp.message)) diff --git a/mediator.py b/mediator.py index 780a141c..0fa4e336 100644 --- a/mediator.py +++ b/mediator.py @@ -1,4 +1,4 @@ -from langchain.memory import PostgresChatMessageHistory +# from langchain.memory import PostgresChatMessageHistory from langchain.schema import Document from langchain.schema.messages import BaseMessage, _message_to_dict, messages_from_dict from langchain.vectorstores import SupabaseVectorStore @@ -10,7 +10,7 @@ import os from dotenv import load_dotenv # Supabase for Postgres Management from supabase.client import create_client, Client -from typing import List +from typing import List, Dict import json load_dotenv() @@ -21,7 +21,7 @@ class SupabaseMediator: def __init__(self): self.supabase: Client = create_client(os.environ['SUPABASE_URL'], os.environ['SUPABASE_KEY']) self.memory_table = os.environ["MEMORY_TABLE"] - self.conversation_table = os.environ["CONVERSATION_TABLE"] + self.session_table = os.environ["SESSION_TABLE"] self.match_function = os.environ["MATCH_FUNCTION"] embeddings = OpenAIEmbeddings( @@ -36,13 +36,83 @@ class SupabaseMediator: table_name=os.environ["VECTOR_TABLE"], query_name=self.match_function ) - # # seed the vector store with facts about bloom - # seed_docs = [ - # Document(page_content="Bloom is your learning companion"), - # Document(page_content="Bloom can be used for learning just about anything! It's your ultimate school assistant."), - # ] - # self.vector_table.add_documents(seed_docs) + # @sentry_sdk.trace + # def conversations(self, location_id: str, user_id: str) -> str | None: + # response = self.supabase.table(self.session_table).select("id").eq("location_id", location_id).eq("user_id", user_id).eq("isActive", True).maybe_single().execute() + # if response: + # session_id = response.data["id"] + # return session_id + # return None + + + + # CRUD for sessions + + @sentry_sdk.trace + def add_session(self, location_id: str, user_id: str, metadata: Dict = {}) -> Dict: + session_id = str(uuid.uuid4()) + payload = { + "id": session_id, + "user_id": user_id, + "location_id": location_id, + "metadata": metadata, + } + representation = self.supabase.table(self.session_table).insert(payload, returning="representation").execute() # type: ignore + print("========================================") + print(representation) + print("========================================") + return representation.data[0] + # self.supabase.table(self.session_table).insert({"id": session_id, "user_id": user_id, "location_id": location_id}).execute() + + @sentry_sdk.trace + def session(self, session_id: str) -> Dict | None: + response = self.supabase.table(self.session_table).select("*").eq("id", session_id).eq("isActive", True).maybe_single().execute() + if response: + return response.data + return None + + @sentry_sdk.trace + def sessions(self, location_id: str, user_id: str, single: bool = True) -> List[Dict] | None: + try: + response = self.supabase.table(self.session_table).select(*["id", "metadata"], count="exact").eq("location_id", location_id).eq("user_id", user_id).eq("isActive", True).order("created_at", desc=True).execute() + if response is not None and response.count is not None: + if (response.count > 1) and single: + # If there is more than 1 active session mark the rest for deletion + session_ids = [record["id"] for record in response.data[1:]] + self._cleanup_sessions(session_ids) # type: ignore + return [response.data[0]] + else: + return response.data + return None + except Exception as e: + print("========================================") + print(e) + print("========================================") + return None + + @sentry_sdk.trace + def update_session(self, user_id: str, session_id: str, metadata: Dict) -> None: + cur = self.supabase.table(self.session_table).select("metadata").eq("id", session_id).eq("user_id", user_id).single().execute() + if cur.data['metadata'] is not None: + new_metadata = cur.data['metadata'].copy() + new_metadata.update(metadata) + else: + new_metadata = metadata + self.supabase.table(self.session_table).update({"metadata": new_metadata}, returning="representation").eq("id", session_id).execute() # type: ignore + + @sentry_sdk.trace + def delete_session(self, user_id: str, session_id: str) -> None: + self.supabase.table(self.session_table).update({"isActive": False}).eq("id", session_id).eq("user_id", user_id).execute() + + # Session Helper Methods + + @sentry_sdk.trace + def _cleanup_sessions(self, session_ids: List[str]) -> None: + for session_id in session_ids: + self.supabase.table(self.session_table).update({"isActive": False}).eq("id", session_id).execute() + + # Session Usage Read & Write @sentry_sdk.trace def messages(self, session_id: str, user_id: str, message_type: str) -> List[BaseMessage]: # type: ignore @@ -53,23 +123,10 @@ class SupabaseMediator: @sentry_sdk.trace def add_message(self, session_id: str, user_id: str, message_type: str, message: BaseMessage) -> None: - self.supabase.table(self.memory_table).insert({"session_id": session_id, "user_id": user_id, "message_type": message_type, "message": _message_to_dict(message)}).execute() - - @sentry_sdk.trace - def conversations(self, location_id: str, user_id: str) -> str | None: - response = self.supabase.table(self.conversation_table).select("id").eq("location_id", location_id).eq("user_id", user_id).eq("isActive", True).maybe_single().execute() - if response: - conversation_id = response.data["id"] - return conversation_id - return None - - @sentry_sdk.trace - def add_conversation(self, location_id: str, user_id: str) -> str: - conversation_id = str(uuid.uuid4()) - self.supabase.table(self.conversation_table).insert({"id": conversation_id, "user_id": user_id, "location_id": location_id}).execute() - return conversation_id - - @sentry_sdk.trace - def delete_conversation(self, conversation_id: str) -> None: - self.supabase.table(self.conversation_table).update({"isActive": False}).eq("id", conversation_id).execute() - + payload = { + "session_id": session_id, + "user_id": user_id, + "message_type": message_type, + "message": _message_to_dict(message) + } + self.supabase.table(self.memory_table).insert(payload).execute()