Skip to content

Commit 462c209

Browse files
apkostkaapkostka
committed
fix: further refactor
1 parent f0f8226 commit 462c209

File tree

2 files changed

+87
-82
lines changed

2 files changed

+87
-82
lines changed

‎backend/tenantfirstaid/chat.py

Lines changed: 67 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
from flask.views import View
44
import os
55

6+
from backend.tenantfirstaid.session import TenantSessionData
7+
68
API_KEY = os.getenv("OPENAI_API_KEY", os.getenv("GITHUB_API_KEY"))
79
BASE_URL = os.getenv("MODEL_ENDPOINT", "https://api.openai.com/v1")
810
MODEL = os.getenv("MODEL_NAME", "o3")
@@ -38,6 +40,66 @@ def __init__(self):
3840
def get_client(self):
3941
return self.client
4042

43+
def prepare_developer_instructions(self, current_session: TenantSessionData):
44+
# Add city and state filters if they are set
45+
instructions = DEFAULT_INSTRUCTIONS
46+
instructions += f"\nThe user is in {current_session['city']} {current_session['state'].upper()}.\n"
47+
return instructions
48+
49+
def prepare_openai_tools(self, current_session: TenantSessionData):
50+
VECTOR_STORE_ID = os.getenv("VECTOR_STORE_ID")
51+
if not VECTOR_STORE_ID:
52+
return None
53+
54+
# We either want to use both city and state, or just state.
55+
# This filters out other cities in the same state.
56+
# The user is gated into selecting a city in Oregon so we don't worry about
57+
# whether the relevant documents exist or not.
58+
return [
59+
{
60+
"type": "file_search",
61+
"vector_store_ids": [VECTOR_STORE_ID],
62+
"max_num_results": os.getenv("NUM_FILE_SEARCH_RESULTS", 10),
63+
"filters": {
64+
"type": "or",
65+
"filters": [
66+
{
67+
"type": "and",
68+
"filters": [
69+
{
70+
"type": "eq",
71+
"key": "city",
72+
"value": current_session["city"],
73+
},
74+
{
75+
"type": "eq",
76+
"key": "state",
77+
"value": current_session["state"],
78+
},
79+
],
80+
}
81+
if current_session["city"] != "null"
82+
else None,
83+
{
84+
"type": "and",
85+
"filters": [
86+
{
87+
"type": "eq",
88+
"key": "city",
89+
"value": "null",
90+
},
91+
{
92+
"type": "eq",
93+
"key": "state",
94+
"value": current_session["state"],
95+
},
96+
],
97+
},
98+
],
99+
},
100+
}
101+
]
102+
41103

42104
class ChatView(View):
43105
client = OpenAI(
@@ -47,94 +109,26 @@ class ChatView(View):
47109

48110
def __init__(self, tenant_session):
49111
self.tenant_session = tenant_session
112+
self.chat_manager = ChatManager()
50113

51114
def dispatch_request(self):
52115
data = request.json
53116
user_msg = data["message"]
54117

55-
# Get or create session ID using Flask sessions
56-
session_id = self.tenant_session.get_flask_session_id()
57-
58118
current_session = self.tenant_session.get()
59119

60-
# Format messages for the new Responses API
61-
input_messages = []
62-
63-
# Add conversation history (excluding system prompt)
64-
for msg in current_session["messages"][0:]:
65-
input_messages.append({"role": msg["role"], "content": msg["content"]})
66-
67-
# Add current user message
68-
input_messages.append({"role": "user", "content": user_msg})
69-
70120
# Update our cache with the user message
71121
current_session["messages"].append({"role": "user", "content": user_msg})
72122

73-
# Add city and state filters if they are set
74-
instructions = DEFAULT_INSTRUCTIONS
75-
instructions += f"\nThe user is in {current_session['city']} {current_session['state'].upper()}.\n"
76-
77-
# We either want to use both city and state, or just state.
78-
# This filters out other cities in the same state.
79-
# The user is gated into selecting a city in Oregon so we don't worry about
80-
# whether the relevant documents exist or not.
81-
VECTOR_STORE_ID = os.getenv("VECTOR_STORE_ID")
82-
83-
tools = (
84-
[
85-
{
86-
"type": "file_search",
87-
"vector_store_ids": [VECTOR_STORE_ID],
88-
"max_num_results": os.getenv("NUM_FILE_SEARCH_RESULTS", 5),
89-
"filters": {
90-
"type": "or",
91-
"filters": [
92-
{
93-
"type": "and",
94-
"filters": [
95-
{
96-
"type": "eq",
97-
"key": "city",
98-
"value": current_session["city"],
99-
},
100-
{
101-
"type": "eq",
102-
"key": "state",
103-
"value": current_session["state"],
104-
},
105-
],
106-
}
107-
if current_session["city"] != "null"
108-
else None,
109-
{
110-
"type": "and",
111-
"filters": [
112-
{
113-
"type": "eq",
114-
"key": "city",
115-
"value": "null",
116-
},
117-
{
118-
"type": "eq",
119-
"key": "state",
120-
"value": current_session["state"],
121-
},
122-
],
123-
},
124-
],
125-
},
126-
}
127-
]
128-
if VECTOR_STORE_ID
129-
else None
130-
)
123+
instructions = self.chat_manager.prepare_developer_instructions(current_session)
124+
tools = self.chat_manager.prepare_openai_tools(current_session)
131125

132126
def generate():
133127
try:
134128
# Use the new Responses API with streaming
135129
response_stream = self.client.responses.create(
136130
model=MODEL,
137-
input=input_messages,
131+
input=current_session["messages"],
138132
instructions=instructions,
139133
reasoning={"effort": MODEL_REASONING_EFFORT},
140134
stream=True,
@@ -165,7 +159,7 @@ def generate():
165159
yield f"Error: {str(e)}"
166160

167161
finally:
168-
self.tenant_session.set(session_id, current_session)
162+
self.tenant_session.set(current_session)
169163

170164
return Response(
171165
stream_with_context(generate()),

‎backend/tenantfirstaid/session.py

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,22 @@
22
import uuid
33
from flask import Response, after_this_request, request, session
44
from flask.views import View
5+
from typing import TypedDict
56
from valkey import Valkey
67
import simplejson as json
78

89

10+
class TenantSessionMessage(TypedDict):
11+
role: str # 'user' or 'assistant'
12+
content: str # The content of the message
13+
14+
15+
class TenantSessionData(TypedDict):
16+
city: str
17+
state: str
18+
messages: list[TenantSessionMessage] # List of messages with role and content
19+
20+
921
# The class to manage tenant sessions using Valkey and Flask sessions
1022
class TenantSession:
1123
def __init__(self):
@@ -30,7 +42,7 @@ def __init__(self):
3042
print(e)
3143

3244
# Retrieves the session ID from Flask session or creates a new one
33-
def get_flask_session_id(self):
45+
def get_flask_session_id(self) -> str:
3446
session_id = session.get("session_id")
3547
if not session_id:
3648
session_id = str(uuid.uuid4())
@@ -43,21 +55,20 @@ def save_session(response):
4355

4456
return session_id
4557

46-
def get(self):
47-
session_id = session.get("session_id")
48-
if not session_id:
49-
return self.getNewSessionData()
58+
def get(self) -> TenantSessionData:
59+
session_id = self.get_flask_session_id()
5060

5161
saved_session = self.db_con.get(session_id)
5262
if not saved_session:
5363
return self.getNewSessionData()
5464

5565
return json.loads(saved_session)
5666

57-
def set(self, session_id, value):
67+
def set(self, value: TenantSessionData):
68+
session_id = self.get_flask_session_id()
5869
self.db_con.set(session_id, json.dumps(value))
5970

60-
def getNewSessionData(self):
71+
def getNewSessionData(self) -> TenantSessionData:
6172
return {
6273
"city": "",
6374
"state": "",
@@ -78,8 +89,8 @@ def dispatch_request(self):
7889
state = data["state"]
7990

8091
# Initialize the session with city and state
81-
initial_data = {"city": city, "state": state, "messages": []}
82-
self.tenant_session.set(session_id, initial_data)
92+
initial_data: TenantSessionData = {"city": city, "state": state, "messages": []}
93+
self.tenant_session.set(initial_data)
8394

8495
return Response(
8596
status=200,

0 commit comments

Comments
 (0)