RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update mediator and cache and add routes

This commit is contained in:
Vineeth Voruganti 2023-12-07 22:36:56 -08:00
parent 4d6a176b25
commit 31e8c9628b
4 changed files with 185 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.gitignore vendored
View File

@ -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

24
cache.py
View File

@ -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

91
main.py
View File

@ -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))

115
mediator.py
View File

@ -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()