Spaces:
Running
Running
File size: 3,862 Bytes
e9d730a d161383 e9d730a d161383 e9d730a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
# src/db/mongodb_store.py
from motor.motor_asyncio import AsyncIOMotorClient
from datetime import datetime
import json
from typing import List, Dict, Optional, Any
from bson import ObjectId
class MongoDBStore:
def __init__(self, mongo_uri: str = "mongodb://localhost:27017"):
"""Initialize MongoDB connection"""
self.client = AsyncIOMotorClient(mongo_uri)
self.db = self.client.rag_chatbot
self.chat_history = self.db.chat_history
self.documents = self.db.documents # Collection for original documents
async def store_document(
self,
document_id: str,
filename: str,
content: str,
content_type: str,
file_size: int
) -> str:
"""Store original document in MongoDB"""
document = {
"document_id": document_id,
"filename": filename,
"content": content,
"content_type": content_type,
"file_size": file_size,
"upload_timestamp": datetime.now()
}
await self.documents.insert_one(document)
return document_id
async def get_document(self, document_id: str) -> Optional[Dict]:
"""Retrieve document by ID"""
return await self.documents.find_one(
{"document_id": document_id},
{"_id": 0} # Exclude MongoDB's _id
)
async def get_all_documents(self) -> List[Dict]:
"""Retrieve all documents"""
cursor = self.documents.find({}, {"_id": 0})
return await cursor.to_list(length=None)
async def store_message(
self,
conversation_id: str,
query: str,
response: str,
context: List[str],
sources: List[Dict],
llm_provider: str
) -> str:
"""Store chat message in MongoDB"""
document = {
"conversation_id": conversation_id,
"timestamp": datetime.now(),
"query": query,
"response": response,
"context": context,
"sources": sources,
"llm_provider": llm_provider,
"feedback": None,
"rating": None
}
result = await self.chat_history.insert_one(document)
return str(result.inserted_id)
async def get_conversation_history(self, conversation_id: str) -> List[Dict]:
"""Retrieve conversation history"""
cursor = self.chat_history.find(
{"conversation_id": conversation_id}
).sort("timestamp", 1)
history = []
async for document in cursor:
document["_id"] = str(document["_id"])
history.append(document)
return history
async def update_feedback(
self,
conversation_id: str,
feedback: Optional[str],
rating: Optional[int]
) -> bool:
"""Update feedback for a conversation"""
result = await self.chat_history.update_many(
{"conversation_id": conversation_id},
{
"$set": {
"feedback": feedback,
"rating": rating
}
}
)
return result.modified_count > 0
async def get_messages_for_summary(self, conversation_id: str) -> List[Dict]:
"""Get messages in format suitable for summarization"""
cursor = self.chat_history.find(
{"conversation_id": conversation_id}
).sort("timestamp", 1)
messages = []
async for doc in cursor:
messages.append({
'role': 'user' if doc['query'] else 'assistant',
'content': doc['query'] or doc['response'],
'timestamp': doc['timestamp'],
'sources': doc['sources']
})
return messages |