Spaces:
Sleeping
Sleeping
Commit
·
9f7b904
1
Parent(s):
d25fc26
fix: improve view last message
Browse files- api/router/bot_one.py +33 -13
- core/chat/chatstore.py +28 -10
api/router/bot_one.py
CHANGED
@@ -20,9 +20,7 @@ from langfuse.llama_index import LlamaIndexCallbackHandler
|
|
20 |
|
21 |
from api.auth import check_user_authentication
|
22 |
from api.router.user import user_dependency
|
23 |
-
from api.function import
|
24 |
-
generate_streaming_completion
|
25 |
-
)
|
26 |
|
27 |
|
28 |
router = APIRouter(tags=["Bot_One"])
|
@@ -35,14 +33,19 @@ def get_chat_store():
|
|
35 |
|
36 |
|
37 |
@router.post("/bot_one/{metadata_id}")
|
38 |
-
async def create_bot_one(
|
|
|
|
|
|
|
|
|
|
|
39 |
auth_response = check_user_authentication(user)
|
40 |
if auth_response:
|
41 |
return auth_response
|
42 |
# Generate a new session ID (UUID)
|
43 |
try:
|
44 |
session_id = generate_uuid()
|
45 |
-
|
46 |
bot_query = BotQuery(user)
|
47 |
bot_query.add_bot(db, session_id, bot_name.name, metadata_id)
|
48 |
|
@@ -93,19 +96,23 @@ async def generator_bot(
|
|
93 |
)
|
94 |
)
|
95 |
else:
|
96 |
-
bot_service = ChatCompletionService(
|
|
|
|
|
97 |
response, metadata, scores = bot_service.generate_completion()
|
98 |
|
99 |
# Set Jakarta timezone
|
100 |
-
jakarta_tz = pytz.timezone(
|
101 |
|
102 |
existing_session = (
|
103 |
-
db.query(Session_Publisher)
|
|
|
|
|
104 |
)
|
105 |
-
|
106 |
existing_session.updated_at = datetime.now(jakarta_tz)
|
107 |
db.commit()
|
108 |
-
|
109 |
return BotResponse(
|
110 |
content=response,
|
111 |
metadata=metadata,
|
@@ -124,12 +131,25 @@ async def get_all_session_bot_one(
|
|
124 |
try:
|
125 |
# Query the session IDs based on the user ID
|
126 |
bot_query = BotQuery(user)
|
|
|
127 |
sessions = bot_query.get_session_ids_bot(db, metadata_id)
|
128 |
-
|
129 |
-
session_data = [
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
130 |
|
131 |
# Convert list of tuples to a simple list
|
132 |
-
session_sorted_data = sorted(
|
|
|
|
|
|
|
|
|
133 |
|
134 |
return session_sorted_data
|
135 |
|
|
|
20 |
|
21 |
from api.auth import check_user_authentication
|
22 |
from api.router.user import user_dependency
|
23 |
+
from api.function import generate_streaming_completion
|
|
|
|
|
24 |
|
25 |
|
26 |
router = APIRouter(tags=["Bot_One"])
|
|
|
33 |
|
34 |
|
35 |
@router.post("/bot_one/{metadata_id}")
|
36 |
+
async def create_bot_one(
|
37 |
+
user: user_dependency,
|
38 |
+
db: db_dependency,
|
39 |
+
metadata_id: int,
|
40 |
+
bot_name: BotCreateRequest,
|
41 |
+
):
|
42 |
auth_response = check_user_authentication(user)
|
43 |
if auth_response:
|
44 |
return auth_response
|
45 |
# Generate a new session ID (UUID)
|
46 |
try:
|
47 |
session_id = generate_uuid()
|
48 |
+
|
49 |
bot_query = BotQuery(user)
|
50 |
bot_query.add_bot(db, session_id, bot_name.name, metadata_id)
|
51 |
|
|
|
96 |
)
|
97 |
)
|
98 |
else:
|
99 |
+
bot_service = ChatCompletionService(
|
100 |
+
session_id, user_prompt_request.prompt, titles, type_bot="specific"
|
101 |
+
)
|
102 |
response, metadata, scores = bot_service.generate_completion()
|
103 |
|
104 |
# Set Jakarta timezone
|
105 |
+
jakarta_tz = pytz.timezone("Asia/Jakarta")
|
106 |
|
107 |
existing_session = (
|
108 |
+
db.query(Session_Publisher)
|
109 |
+
.filter(Session_Publisher.id == session_id)
|
110 |
+
.first()
|
111 |
)
|
112 |
+
|
113 |
existing_session.updated_at = datetime.now(jakarta_tz)
|
114 |
db.commit()
|
115 |
+
|
116 |
return BotResponse(
|
117 |
content=response,
|
118 |
metadata=metadata,
|
|
|
131 |
try:
|
132 |
# Query the session IDs based on the user ID
|
133 |
bot_query = BotQuery(user)
|
134 |
+
chat_store = ChatStore()
|
135 |
sessions = bot_query.get_session_ids_bot(db, metadata_id)
|
136 |
+
|
137 |
+
session_data = [
|
138 |
+
{
|
139 |
+
"id": session.id,
|
140 |
+
"bot_name": session.bot_name,
|
141 |
+
"updated_at": str(session.updated_at),
|
142 |
+
"last_message": chat_store.get_last_message_mongodb(session.id),
|
143 |
+
}
|
144 |
+
for session in sessions
|
145 |
+
]
|
146 |
|
147 |
# Convert list of tuples to a simple list
|
148 |
+
session_sorted_data = sorted(
|
149 |
+
session_data,
|
150 |
+
key=lambda x: datetime.fromisoformat(x["updated_at"]),
|
151 |
+
reverse=True,
|
152 |
+
)
|
153 |
|
154 |
return session_sorted_data
|
155 |
|
core/chat/chatstore.py
CHANGED
@@ -21,7 +21,7 @@ class ChatStore:
|
|
21 |
self.client = MongoClient(uri)
|
22 |
|
23 |
def initialize_memory_bot(self, session_id):
|
24 |
-
|
25 |
chat_store = RedisChatStore(
|
26 |
redis_client=self.redis_client, ttl=86400 # Time-to-live set for 1 hour
|
27 |
)
|
@@ -56,17 +56,28 @@ class ChatStore:
|
|
56 |
|
57 |
# Decode and parse each item into a dictionary
|
58 |
return [json.loads(m.decode("utf-8")) for m in items]
|
59 |
-
|
60 |
def get_last_message(self, session_id: str) -> Optional[Dict]:
|
61 |
"""Get the last message for a session_id."""
|
62 |
last_message = self.redis_client.lindex(session_id, -1)
|
63 |
-
|
64 |
if last_message is None:
|
65 |
return None # Return None if there are no messages
|
66 |
-
|
67 |
# Decode and parse the last message into a dictionary
|
68 |
return json.loads(last_message.decode("utf-8"))
|
69 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
70 |
|
71 |
def delete_last_message(self, session_id: str) -> Optional[ChatMessage]:
|
72 |
"""Delete last message for a session_id."""
|
@@ -82,18 +93,22 @@ class ChatStore:
|
|
82 |
def clean_message(self, session_id: str) -> Optional[ChatMessage]:
|
83 |
"""Delete specific message for a session_id."""
|
84 |
current_list = self.redis_client.lrange(session_id, 0, -1)
|
85 |
-
|
86 |
indices_to_delete = []
|
87 |
for index, item in enumerate(current_list):
|
88 |
data = json.loads(item) # Parse JSON string to dict
|
89 |
|
90 |
# Logic to determine if item should be removed
|
91 |
-
if (data.get("role") == "assistant" and data.get("content") is None) or (
|
|
|
|
|
92 |
indices_to_delete.append(index)
|
93 |
|
94 |
# Remove elements by their indices in reverse order
|
95 |
for index in reversed(indices_to_delete):
|
96 |
-
self.redis_client.lrem(
|
|
|
|
|
97 |
|
98 |
def get_keys(self) -> List[str]:
|
99 |
"""Get all keys."""
|
@@ -154,7 +169,7 @@ class ChatStore:
|
|
154 |
|
155 |
# Convert the cursor to a list and exclude the _id field
|
156 |
documents_list = [
|
157 |
-
{key: doc[key] for key in doc if key !="_id" and doc[key] is not None}
|
158 |
for doc in documents
|
159 |
]
|
160 |
|
@@ -162,7 +177,10 @@ class ChatStore:
|
|
162 |
print(documents_list) # Optional: If you want to see the output
|
163 |
|
164 |
return documents_list
|
165 |
-
|
166 |
except Exception as e:
|
167 |
print(f"An error occurred while retrieving messages: {e}")
|
168 |
-
return JSONResponse(
|
|
|
|
|
|
|
|
21 |
self.client = MongoClient(uri)
|
22 |
|
23 |
def initialize_memory_bot(self, session_id):
|
24 |
+
|
25 |
chat_store = RedisChatStore(
|
26 |
redis_client=self.redis_client, ttl=86400 # Time-to-live set for 1 hour
|
27 |
)
|
|
|
56 |
|
57 |
# Decode and parse each item into a dictionary
|
58 |
return [json.loads(m.decode("utf-8")) for m in items]
|
59 |
+
|
60 |
def get_last_message(self, session_id: str) -> Optional[Dict]:
|
61 |
"""Get the last message for a session_id."""
|
62 |
last_message = self.redis_client.lindex(session_id, -1)
|
63 |
+
|
64 |
if last_message is None:
|
65 |
return None # Return None if there are no messages
|
66 |
+
|
67 |
# Decode and parse the last message into a dictionary
|
68 |
return json.loads(last_message.decode("utf-8"))
|
69 |
|
70 |
+
def get_last_message_mongodb(self, session_id: str):
|
71 |
+
db = self.client["bot_database"]
|
72 |
+
collection = db[session_id]
|
73 |
+
|
74 |
+
# Get the last document by sorting by _id in descending order
|
75 |
+
last_document = collection.find().sort("_id", -1).limit(1)
|
76 |
+
|
77 |
+
for doc in last_document:
|
78 |
+
doc["content"]
|
79 |
+
|
80 |
+
return str(doc["content"])
|
81 |
|
82 |
def delete_last_message(self, session_id: str) -> Optional[ChatMessage]:
|
83 |
"""Delete last message for a session_id."""
|
|
|
93 |
def clean_message(self, session_id: str) -> Optional[ChatMessage]:
|
94 |
"""Delete specific message for a session_id."""
|
95 |
current_list = self.redis_client.lrange(session_id, 0, -1)
|
96 |
+
|
97 |
indices_to_delete = []
|
98 |
for index, item in enumerate(current_list):
|
99 |
data = json.loads(item) # Parse JSON string to dict
|
100 |
|
101 |
# Logic to determine if item should be removed
|
102 |
+
if (data.get("role") == "assistant" and data.get("content") is None) or (
|
103 |
+
data.get("role") == "tool"
|
104 |
+
):
|
105 |
indices_to_delete.append(index)
|
106 |
|
107 |
# Remove elements by their indices in reverse order
|
108 |
for index in reversed(indices_to_delete):
|
109 |
+
self.redis_client.lrem(
|
110 |
+
session_id, 1, current_list[index]
|
111 |
+
) # Remove the element from the list in Redis
|
112 |
|
113 |
def get_keys(self) -> List[str]:
|
114 |
"""Get all keys."""
|
|
|
169 |
|
170 |
# Convert the cursor to a list and exclude the _id field
|
171 |
documents_list = [
|
172 |
+
{key: doc[key] for key in doc if key != "_id" and doc[key] is not None}
|
173 |
for doc in documents
|
174 |
]
|
175 |
|
|
|
177 |
print(documents_list) # Optional: If you want to see the output
|
178 |
|
179 |
return documents_list
|
180 |
+
|
181 |
except Exception as e:
|
182 |
print(f"An error occurred while retrieving messages: {e}")
|
183 |
+
return JSONResponse(
|
184 |
+
status_code=500,
|
185 |
+
content=f"An error occurred while retrieving messages: {e}",
|
186 |
+
)
|