Spaces:
Sleeping
Sleeping
Chat-engine cache
Browse files- app/api/postgresql_routes.py +44 -266
- app/api/rag_routes.py +469 -6
- app/database/models.py +3 -1
- app/models/rag_models.py +58 -2
- app/utils/cache_config.py +45 -0
- beach_request.json +0 -0
- chat_request.json +0 -0
- pytest.ini +0 -12
- test_body.json +0 -0
- test_rag_api.py +0 -263
- update_body.json +0 -0
app/api/postgresql_routes.py
CHANGED
@@ -19,8 +19,6 @@ from sqlalchemy.exc import SQLAlchemyError
|
|
19 |
from sqlalchemy import desc, func
|
20 |
from cachetools import TTLCache
|
21 |
import uuid
|
22 |
-
import asyncio
|
23 |
-
import httpx # Import httpx for HTTP requests
|
24 |
|
25 |
from app.database.postgresql import get_db
|
26 |
from app.database.models import FAQItem, EmergencyItem, EventItem, AboutPixity, SolanaSummit, DaNangBucketList, ApiKey, VectorDatabase, Document, VectorStatus, TelegramBot, ChatEngine, BotEngine, EngineVectorDb, DocumentContent
|
@@ -3204,7 +3202,6 @@ class VectorStatusBase(BaseModel):
|
|
3204 |
document_id: int
|
3205 |
vector_database_id: int
|
3206 |
vector_id: Optional[str] = None
|
3207 |
-
document_name: Optional[str] = None # Added to match database schema
|
3208 |
status: str = "pending"
|
3209 |
error_message: Optional[str] = None
|
3210 |
|
@@ -3668,167 +3665,69 @@ async def update_document(
|
|
3668 |
db.add(document_content)
|
3669 |
|
3670 |
# Get vector status for Pinecone cleanup
|
3671 |
-
vector_status = db.query(VectorStatus).filter(
|
3672 |
-
VectorStatus.document_id == document_id,
|
3673 |
-
VectorStatus.vector_database_id == document.vector_database_id
|
3674 |
-
).first()
|
3675 |
|
3676 |
# Store old vector_id for cleanup
|
3677 |
old_vector_id = None
|
3678 |
if vector_status and vector_status.vector_id:
|
3679 |
old_vector_id = vector_status.vector_id
|
3680 |
-
logger.info(f"Found old vector_id {old_vector_id} for document {document_id}, planning to delete")
|
3681 |
|
3682 |
-
#
|
3683 |
if vector_status:
|
3684 |
-
|
3685 |
-
|
3686 |
-
|
3687 |
-
|
3688 |
-
|
3689 |
-
|
3690 |
-
# Create new vector status
|
3691 |
vector_status = VectorStatus(
|
3692 |
document_id=document_id,
|
3693 |
vector_database_id=document.vector_database_id,
|
3694 |
-
|
3695 |
-
document_name=document.name
|
3696 |
)
|
3697 |
db.add(vector_status)
|
3698 |
-
db.flush()
|
3699 |
-
logger.info(f"Created new vector status for document {document_id} with status 'pending'")
|
3700 |
|
3701 |
-
#
|
3702 |
-
if old_vector_id and vector_db and document.vector_database_id:
|
3703 |
try:
|
3704 |
-
|
3705 |
-
|
3706 |
-
# Call PDF API to delete document using HTTP request (avoids circular imports)
|
3707 |
-
base_url = "http://localhost:8000"
|
3708 |
-
delete_url = f"{base_url}/pdf/document"
|
3709 |
|
3710 |
-
|
3711 |
-
|
3712 |
-
|
3713 |
-
|
3714 |
-
|
3715 |
-
}
|
3716 |
|
3717 |
-
|
|
|
|
|
|
|
|
|
3718 |
|
3719 |
-
|
3720 |
-
async with httpx.AsyncClient() as client:
|
3721 |
-
response = await client.delete(delete_url, params=params)
|
3722 |
-
if response.status_code == 200:
|
3723 |
-
result = response.json()
|
3724 |
-
vectors_deleted = result.get('vectors_deleted', 0)
|
3725 |
-
logger.info(f"Successfully deleted {vectors_deleted} old vectors for document {document_id}")
|
3726 |
-
else:
|
3727 |
-
logger.warning(f"Failed to delete old vectors: {response.status_code} - {response.text}")
|
3728 |
except Exception as e:
|
3729 |
-
logger.error(f"Error
|
3730 |
-
# Continue with the update even if vector deletion fails
|
3731 |
|
3732 |
-
#
|
3733 |
-
if document.vector_database_id:
|
3734 |
try:
|
3735 |
-
#
|
3736 |
-
|
3737 |
-
import tempfile
|
3738 |
-
import os
|
3739 |
|
3740 |
-
|
3741 |
-
|
3742 |
-
|
3743 |
-
|
3744 |
-
|
3745 |
-
|
3746 |
-
with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_extension}") as temp_file:
|
3747 |
-
temp_file.write(file_content)
|
3748 |
-
temp_path = temp_file.name
|
3749 |
-
|
3750 |
-
# Prepare multipart form data for the upload
|
3751 |
-
import aiofiles
|
3752 |
-
from aiofiles import os as aio_os
|
3753 |
-
|
3754 |
-
async with httpx.AsyncClient(timeout=300) as client: # Increased timeout to 300 seconds
|
3755 |
-
# Open the temp file for async reading
|
3756 |
-
async with aiofiles.open(temp_path, "rb") as f:
|
3757 |
-
file_data = await f.read()
|
3758 |
-
|
3759 |
-
# Create form data
|
3760 |
-
files = {"file": (filename, file_data, document.content_type)}
|
3761 |
-
form_data = {
|
3762 |
-
"title": document.name,
|
3763 |
-
"vector_database_id": str(document.vector_database_id),
|
3764 |
-
"namespace": f"vdb-{document.vector_database_id}"
|
3765 |
-
}
|
3766 |
-
|
3767 |
-
# Call the PDF upload API
|
3768 |
-
upload_url = f"{base_url}/pdf/upload"
|
3769 |
-
logger.info(f"Calling PDF upload API for document {document_id}")
|
3770 |
-
|
3771 |
-
response = await client.post(upload_url, files=files, data=form_data)
|
3772 |
-
|
3773 |
-
# Process the response
|
3774 |
-
if response.status_code == 200:
|
3775 |
-
result = response.json()
|
3776 |
-
logger.info(f"Successfully uploaded document {document_id}: {result}")
|
3777 |
-
|
3778 |
-
# Get the new vector_id from the result
|
3779 |
-
new_vector_id = result.get('document_id')
|
3780 |
-
|
3781 |
-
# If upload was successful, update the vector status in PostgreSQL
|
3782 |
-
if result.get('success') and new_vector_id:
|
3783 |
-
# Get the latest vector status
|
3784 |
-
await asyncio.sleep(1) # Small delay to ensure DB consistency
|
3785 |
-
|
3786 |
-
# Use a new DB session for this update
|
3787 |
-
from app.database.postgresql import SessionLocal
|
3788 |
-
async_db = SessionLocal()
|
3789 |
-
try:
|
3790 |
-
vs = async_db.query(VectorStatus).filter(
|
3791 |
-
VectorStatus.document_id == document_id,
|
3792 |
-
VectorStatus.vector_database_id == document.vector_database_id
|
3793 |
-
).first()
|
3794 |
-
|
3795 |
-
if vs:
|
3796 |
-
vs.vector_id = new_vector_id
|
3797 |
-
vs.status = "completed"
|
3798 |
-
vs.embedded_at = datetime.now()
|
3799 |
-
|
3800 |
-
# Also update document embedded status
|
3801 |
-
doc = async_db.query(Document).filter(Document.id == document_id).first()
|
3802 |
-
if doc:
|
3803 |
-
doc.is_embedded = True
|
3804 |
-
|
3805 |
-
async_db.commit()
|
3806 |
-
logger.info(f"Updated vector status with new vector_id {new_vector_id}")
|
3807 |
-
finally:
|
3808 |
-
async_db.close()
|
3809 |
-
else:
|
3810 |
-
logger.error(f"Failed to upload document: {response.status_code} - {response.text}")
|
3811 |
-
|
3812 |
-
# Clean up temporary file
|
3813 |
-
try:
|
3814 |
-
await aio_os.remove(temp_path)
|
3815 |
-
except Exception as cleanup_error:
|
3816 |
-
logger.error(f"Error cleaning up temporary file: {str(cleanup_error)}")
|
3817 |
-
except Exception as e:
|
3818 |
-
logger.error(f"Error in background upload task: {str(e)}")
|
3819 |
-
logger.error(traceback.format_exc())
|
3820 |
|
3821 |
-
|
3822 |
-
if background_tasks:
|
3823 |
-
background_tasks.add_task(upload_and_process_document)
|
3824 |
-
logger.info("Added document upload to background tasks")
|
3825 |
-
else:
|
3826 |
-
logger.warning("Background tasks not available, skipping document upload")
|
3827 |
except Exception as e:
|
3828 |
-
logger.error(f"Error
|
3829 |
-
# Continue with the update even if
|
3830 |
|
3831 |
-
# Commit changes to document record
|
3832 |
db.commit()
|
3833 |
db.refresh(document)
|
3834 |
|
@@ -3878,152 +3777,31 @@ async def delete_document(
|
|
3878 |
- **document_id**: ID of the document to delete
|
3879 |
"""
|
3880 |
try:
|
3881 |
-
logger.info(f"Starting deletion process for document ID {document_id}")
|
3882 |
-
|
3883 |
# Check if document exists
|
3884 |
document = db.query(Document).filter(Document.id == document_id).first()
|
3885 |
if not document:
|
3886 |
-
logger.warning(f"Document with ID {document_id} not found for deletion")
|
3887 |
raise HTTPException(status_code=404, detail=f"Document with ID {document_id} not found")
|
3888 |
|
3889 |
-
vector_database_id = document.vector_database_id
|
3890 |
-
logger.info(f"Found document to delete: name={document.name}, vector_database_id={vector_database_id}")
|
3891 |
-
|
3892 |
-
# Get the vector_id from VectorStatus before deletion
|
3893 |
-
vector_status = db.query(VectorStatus).filter(
|
3894 |
-
VectorStatus.document_id == document_id,
|
3895 |
-
VectorStatus.vector_database_id == vector_database_id
|
3896 |
-
).first()
|
3897 |
-
|
3898 |
-
# Store the vector_id for Pinecone deletion
|
3899 |
-
vector_id = None
|
3900 |
-
pinecone_deletion_success = False
|
3901 |
-
pinecone_error = None
|
3902 |
-
|
3903 |
-
if vector_status and vector_status.vector_id:
|
3904 |
-
vector_id = vector_status.vector_id
|
3905 |
-
logger.info(f"Found vector_id {vector_id} for document {document_id}")
|
3906 |
-
|
3907 |
-
# Get vector database info
|
3908 |
-
vector_db = db.query(VectorDatabase).filter(
|
3909 |
-
VectorDatabase.id == vector_database_id
|
3910 |
-
).first()
|
3911 |
-
|
3912 |
-
if vector_db:
|
3913 |
-
logger.info(f"Found vector database: name={vector_db.name}, index={vector_db.pinecone_index}")
|
3914 |
-
|
3915 |
-
# Create namespace for vector database
|
3916 |
-
namespace = f"vdb-{vector_database_id}"
|
3917 |
-
|
3918 |
-
try:
|
3919 |
-
import httpx
|
3920 |
-
|
3921 |
-
# Call PDF API to delete from Pinecone using an HTTP request
|
3922 |
-
# This avoids circular import issues
|
3923 |
-
base_url = "http://localhost:8000" # Adjust this to match your actual base URL
|
3924 |
-
delete_url = f"{base_url}/pdf/document"
|
3925 |
-
|
3926 |
-
params = {
|
3927 |
-
"document_id": vector_id, # Use the vector_id instead of the PostgreSQL document ID
|
3928 |
-
"namespace": namespace,
|
3929 |
-
"index_name": vector_db.pinecone_index,
|
3930 |
-
"vector_database_id": vector_database_id
|
3931 |
-
}
|
3932 |
-
|
3933 |
-
logger.info(f"Calling PDF API to delete vectors with params: {params}")
|
3934 |
-
|
3935 |
-
# Add retry logic for better reliability
|
3936 |
-
max_retries = 3
|
3937 |
-
retry_delay = 2 # seconds
|
3938 |
-
success = False
|
3939 |
-
last_error = None
|
3940 |
-
|
3941 |
-
for retry in range(max_retries):
|
3942 |
-
try:
|
3943 |
-
async with httpx.AsyncClient(timeout=300) as client: # Increased timeout to 300 seconds
|
3944 |
-
response = await client.delete(delete_url, params=params)
|
3945 |
-
|
3946 |
-
if response.status_code == 200:
|
3947 |
-
result = response.json()
|
3948 |
-
pinecone_deletion_success = result.get('success', False)
|
3949 |
-
vectors_deleted = result.get('vectors_deleted', 0)
|
3950 |
-
logger.info(f"Vector deletion API call response: success={pinecone_deletion_success}, vectors_deleted={vectors_deleted}")
|
3951 |
-
success = True
|
3952 |
-
break
|
3953 |
-
else:
|
3954 |
-
last_error = f"Failed with status code {response.status_code}: {response.text}"
|
3955 |
-
logger.warning(f"Deletion attempt {retry+1}/{max_retries} failed: {last_error}")
|
3956 |
-
except Exception as e:
|
3957 |
-
last_error = str(e)
|
3958 |
-
logger.warning(f"Deletion attempt {retry+1}/{max_retries} failed with exception: {last_error}")
|
3959 |
-
|
3960 |
-
# Wait before retrying
|
3961 |
-
if retry < max_retries - 1: # Don't sleep after the last attempt
|
3962 |
-
logger.info(f"Retrying in {retry_delay} seconds...")
|
3963 |
-
await asyncio.sleep(retry_delay)
|
3964 |
-
retry_delay *= 2 # Exponential backoff
|
3965 |
-
|
3966 |
-
if not success:
|
3967 |
-
pinecone_error = f"All deletion attempts failed. Last error: {last_error}"
|
3968 |
-
logger.warning(pinecone_error)
|
3969 |
-
# Continue with PostgreSQL deletion even if Pinecone deletion fails
|
3970 |
-
except Exception as e:
|
3971 |
-
pinecone_error = f"Error setting up Pinecone deletion: {str(e)}"
|
3972 |
-
logger.error(pinecone_error)
|
3973 |
-
# Continue with PostgreSQL deletion even if Pinecone deletion fails
|
3974 |
-
else:
|
3975 |
-
logger.warning(f"No vector_id found for document {document_id}, skipping Pinecone deletion")
|
3976 |
-
|
3977 |
# Delete vector status
|
3978 |
-
|
3979 |
-
logger.info(f"Deleted {result_vs} vector status records for document {document_id}")
|
3980 |
|
3981 |
# Delete document content
|
3982 |
-
|
3983 |
-
logger.info(f"Deleted {result_dc} document content records for document {document_id}")
|
3984 |
|
3985 |
# Delete document
|
3986 |
db.delete(document)
|
3987 |
db.commit()
|
3988 |
-
logger.info(f"Document with ID {document_id} successfully deleted from PostgreSQL")
|
3989 |
-
|
3990 |
-
# Prepare response with information about what happened
|
3991 |
-
response = {
|
3992 |
-
"status": "success",
|
3993 |
-
"message": f"Document with ID {document_id} deleted successfully",
|
3994 |
-
"postgresql_deletion": {
|
3995 |
-
"document_deleted": True,
|
3996 |
-
"vector_status_deleted": result_vs > 0,
|
3997 |
-
"document_content_deleted": result_dc > 0
|
3998 |
-
}
|
3999 |
-
}
|
4000 |
-
|
4001 |
-
# Add Pinecone deletion information
|
4002 |
-
if vector_id:
|
4003 |
-
response["pinecone_deletion"] = {
|
4004 |
-
"attempted": True,
|
4005 |
-
"vector_id": vector_id,
|
4006 |
-
"success": pinecone_deletion_success,
|
4007 |
-
}
|
4008 |
-
if pinecone_error:
|
4009 |
-
response["pinecone_deletion"]["error"] = pinecone_error
|
4010 |
-
else:
|
4011 |
-
response["pinecone_deletion"] = {
|
4012 |
-
"attempted": False,
|
4013 |
-
"reason": "No vector_id found for document"
|
4014 |
-
}
|
4015 |
|
4016 |
-
return
|
4017 |
except HTTPException:
|
4018 |
-
logger.warning(f"HTTP exception in delete_document for ID {document_id}")
|
4019 |
raise
|
4020 |
except SQLAlchemyError as e:
|
4021 |
db.rollback()
|
4022 |
-
logger.error(f"Database error deleting document
|
4023 |
logger.error(traceback.format_exc())
|
4024 |
raise HTTPException(status_code=500, detail=f"Database error: {str(e)}")
|
4025 |
except Exception as e:
|
4026 |
db.rollback()
|
4027 |
-
logger.error(f"Error deleting document
|
4028 |
logger.error(traceback.format_exc())
|
4029 |
raise HTTPException(status_code=500, detail=f"Error deleting document: {str(e)}")
|
|
|
19 |
from sqlalchemy import desc, func
|
20 |
from cachetools import TTLCache
|
21 |
import uuid
|
|
|
|
|
22 |
|
23 |
from app.database.postgresql import get_db
|
24 |
from app.database.models import FAQItem, EmergencyItem, EventItem, AboutPixity, SolanaSummit, DaNangBucketList, ApiKey, VectorDatabase, Document, VectorStatus, TelegramBot, ChatEngine, BotEngine, EngineVectorDb, DocumentContent
|
|
|
3202 |
document_id: int
|
3203 |
vector_database_id: int
|
3204 |
vector_id: Optional[str] = None
|
|
|
3205 |
status: str = "pending"
|
3206 |
error_message: Optional[str] = None
|
3207 |
|
|
|
3665 |
db.add(document_content)
|
3666 |
|
3667 |
# Get vector status for Pinecone cleanup
|
3668 |
+
vector_status = db.query(VectorStatus).filter(VectorStatus.document_id == document_id).first()
|
|
|
|
|
|
|
3669 |
|
3670 |
# Store old vector_id for cleanup
|
3671 |
old_vector_id = None
|
3672 |
if vector_status and vector_status.vector_id:
|
3673 |
old_vector_id = vector_status.vector_id
|
|
|
3674 |
|
3675 |
+
# Update vector status to pending
|
3676 |
if vector_status:
|
3677 |
+
vector_status.status = "pending"
|
3678 |
+
vector_status.vector_id = None
|
3679 |
+
vector_status.embedded_at = None
|
3680 |
+
vector_status.error_message = None
|
3681 |
+
else:
|
3682 |
+
# Create new vector status if it doesn't exist
|
|
|
3683 |
vector_status = VectorStatus(
|
3684 |
document_id=document_id,
|
3685 |
vector_database_id=document.vector_database_id,
|
3686 |
+
status="pending"
|
|
|
3687 |
)
|
3688 |
db.add(vector_status)
|
|
|
|
|
3689 |
|
3690 |
+
# Schedule deletion of old vectors in Pinecone if we have all needed info
|
3691 |
+
if old_vector_id and vector_db and document.vector_database_id and background_tasks:
|
3692 |
try:
|
3693 |
+
# Initialize PDFProcessor for vector deletion
|
3694 |
+
from app.pdf.processor import PDFProcessor
|
|
|
|
|
|
|
3695 |
|
3696 |
+
processor = PDFProcessor(
|
3697 |
+
index_name=vector_db.pinecone_index,
|
3698 |
+
namespace=f"vdb-{document.vector_database_id}",
|
3699 |
+
vector_db_id=document.vector_database_id
|
3700 |
+
)
|
|
|
3701 |
|
3702 |
+
# Add deletion task to background tasks
|
3703 |
+
background_tasks.add_task(
|
3704 |
+
processor.delete_document_vectors,
|
3705 |
+
old_vector_id
|
3706 |
+
)
|
3707 |
|
3708 |
+
logger.info(f"Scheduled deletion of old vectors for document {document_id}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
3709 |
except Exception as e:
|
3710 |
+
logger.error(f"Error scheduling vector deletion: {str(e)}")
|
3711 |
+
# Continue with the update even if vector deletion scheduling fails
|
3712 |
|
3713 |
+
# Schedule document for re-embedding if possible
|
3714 |
+
if background_tasks and document.vector_database_id:
|
3715 |
try:
|
3716 |
+
# Import here to avoid circular imports
|
3717 |
+
from app.pdf.tasks import process_document_for_embedding
|
|
|
|
|
3718 |
|
3719 |
+
# Schedule embedding
|
3720 |
+
background_tasks.add_task(
|
3721 |
+
process_document_for_embedding,
|
3722 |
+
document_id=document_id,
|
3723 |
+
vector_db_id=document.vector_database_id
|
3724 |
+
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
3725 |
|
3726 |
+
logger.info(f"Scheduled re-embedding for document {document_id}")
|
|
|
|
|
|
|
|
|
|
|
3727 |
except Exception as e:
|
3728 |
+
logger.error(f"Error scheduling document embedding: {str(e)}")
|
3729 |
+
# Continue with the update even if embedding scheduling fails
|
3730 |
|
|
|
3731 |
db.commit()
|
3732 |
db.refresh(document)
|
3733 |
|
|
|
3777 |
- **document_id**: ID of the document to delete
|
3778 |
"""
|
3779 |
try:
|
|
|
|
|
3780 |
# Check if document exists
|
3781 |
document = db.query(Document).filter(Document.id == document_id).first()
|
3782 |
if not document:
|
|
|
3783 |
raise HTTPException(status_code=404, detail=f"Document with ID {document_id} not found")
|
3784 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
3785 |
# Delete vector status
|
3786 |
+
db.query(VectorStatus).filter(VectorStatus.document_id == document_id).delete()
|
|
|
3787 |
|
3788 |
# Delete document content
|
3789 |
+
db.query(DocumentContent).filter(DocumentContent.document_id == document_id).delete()
|
|
|
3790 |
|
3791 |
# Delete document
|
3792 |
db.delete(document)
|
3793 |
db.commit()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
3794 |
|
3795 |
+
return {"status": "success", "message": f"Document with ID {document_id} deleted successfully"}
|
3796 |
except HTTPException:
|
|
|
3797 |
raise
|
3798 |
except SQLAlchemyError as e:
|
3799 |
db.rollback()
|
3800 |
+
logger.error(f"Database error deleting document: {e}")
|
3801 |
logger.error(traceback.format_exc())
|
3802 |
raise HTTPException(status_code=500, detail=f"Database error: {str(e)}")
|
3803 |
except Exception as e:
|
3804 |
db.rollback()
|
3805 |
+
logger.error(f"Error deleting document: {e}")
|
3806 |
logger.error(traceback.format_exc())
|
3807 |
raise HTTPException(status_code=500, detail=f"Error deleting document: {str(e)}")
|
app/api/rag_routes.py
CHANGED
@@ -1,4 +1,4 @@
|
|
1 |
-
from fastapi import APIRouter, HTTPException, Depends, Query, BackgroundTasks, Request
|
2 |
from typing import List, Optional, Dict, Any
|
3 |
import logging
|
4 |
import time
|
@@ -12,8 +12,23 @@ from datetime import datetime
|
|
12 |
from langchain.prompts import PromptTemplate
|
13 |
from langchain_google_genai import GoogleGenerativeAIEmbeddings
|
14 |
from app.utils.utils import timer_decorator
|
|
|
|
|
15 |
|
16 |
from app.database.mongodb import get_chat_history, get_request_history, session_collection
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
17 |
from app.database.pinecone import (
|
18 |
search_vectors,
|
19 |
get_chain,
|
@@ -30,7 +45,12 @@ from app.models.rag_models import (
|
|
30 |
SourceDocument,
|
31 |
EmbeddingRequest,
|
32 |
EmbeddingResponse,
|
33 |
-
UserMessageModel
|
|
|
|
|
|
|
|
|
|
|
34 |
)
|
35 |
|
36 |
# Configure logging
|
@@ -75,15 +95,15 @@ prompt = PromptTemplate(
|
|
75 |
You are Pixity - a professional tour guide assistant that assists users in finding information about places in Da Nang, Vietnam.
|
76 |
You can provide details on restaurants, cafes, hotels, attractions, and other local venues.
|
77 |
You have to use core knowledge and conversation history to chat with users, who are Da Nang's tourists.
|
78 |
-
Pixity
|
79 |
Naturally Cute: Shows cuteness through word choice, soft emojis, and gentle care for the user.
|
80 |
Playful – a little bit cheeky in a lovable way: Occasionally cracks jokes, uses light memes or throws in a surprise response that makes users smile. Think Duolingo-style humor, but less threatening.
|
81 |
Smart & Proactive: Friendly, but also delivers quick, accurate info. Knows how to guide users to the right place – at the right time – with the right solution.
|
82 |
-
Tone & Voice: Friendly – Youthful – Snappy. Uses simple words, similar to daily chat language (e.g.,
|
83 |
SAMPLE DIALOGUES
|
84 |
When a user opens the chatbot for the first time:
|
85 |
User: Hello?
|
86 |
-
Pixity: Hi hi 👋 I
|
87 |
|
88 |
Return Format:
|
89 |
Respond in friendly, natural, concise and use only English like a real tour guide.
|
@@ -344,4 +364,447 @@ async def health_check():
|
|
344 |
"services": services,
|
345 |
"retrieval_config": retrieval_config,
|
346 |
"timestamp": datetime.now().isoformat()
|
347 |
-
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from fastapi import APIRouter, HTTPException, Depends, Query, BackgroundTasks, Request, Path, Body, status
|
2 |
from typing import List, Optional, Dict, Any
|
3 |
import logging
|
4 |
import time
|
|
|
12 |
from langchain.prompts import PromptTemplate
|
13 |
from langchain_google_genai import GoogleGenerativeAIEmbeddings
|
14 |
from app.utils.utils import timer_decorator
|
15 |
+
from sqlalchemy.orm import Session
|
16 |
+
from sqlalchemy.exc import SQLAlchemyError
|
17 |
|
18 |
from app.database.mongodb import get_chat_history, get_request_history, session_collection
|
19 |
+
from app.database.postgresql import get_db
|
20 |
+
from app.database.models import ChatEngine
|
21 |
+
from app.utils.cache import get_cache, InMemoryCache
|
22 |
+
from app.utils.cache_config import (
|
23 |
+
CHAT_ENGINE_CACHE_TTL,
|
24 |
+
MODEL_CONFIG_CACHE_TTL,
|
25 |
+
RETRIEVER_CACHE_TTL,
|
26 |
+
PROMPT_TEMPLATE_CACHE_TTL,
|
27 |
+
get_chat_engine_cache_key,
|
28 |
+
get_model_config_cache_key,
|
29 |
+
get_retriever_cache_key,
|
30 |
+
get_prompt_template_cache_key
|
31 |
+
)
|
32 |
from app.database.pinecone import (
|
33 |
search_vectors,
|
34 |
get_chain,
|
|
|
45 |
SourceDocument,
|
46 |
EmbeddingRequest,
|
47 |
EmbeddingResponse,
|
48 |
+
UserMessageModel,
|
49 |
+
ChatEngineBase,
|
50 |
+
ChatEngineCreate,
|
51 |
+
ChatEngineUpdate,
|
52 |
+
ChatEngineResponse,
|
53 |
+
ChatWithEngineRequest
|
54 |
)
|
55 |
|
56 |
# Configure logging
|
|
|
95 |
You are Pixity - a professional tour guide assistant that assists users in finding information about places in Da Nang, Vietnam.
|
96 |
You can provide details on restaurants, cafes, hotels, attractions, and other local venues.
|
97 |
You have to use core knowledge and conversation history to chat with users, who are Da Nang's tourists.
|
98 |
+
Pixity's Core Personality: Friendly & Warm: Chats like a trustworthy friend who listens and is always ready to help.
|
99 |
Naturally Cute: Shows cuteness through word choice, soft emojis, and gentle care for the user.
|
100 |
Playful – a little bit cheeky in a lovable way: Occasionally cracks jokes, uses light memes or throws in a surprise response that makes users smile. Think Duolingo-style humor, but less threatening.
|
101 |
Smart & Proactive: Friendly, but also delivers quick, accurate info. Knows how to guide users to the right place – at the right time – with the right solution.
|
102 |
+
Tone & Voice: Friendly – Youthful – Snappy. Uses simple words, similar to daily chat language (e.g., "Let's find it together!" / "Need a tip?" / "Here's something cool"). Avoids sounding robotic or overly scripted. Can joke lightly in smart ways, making Pixity feel like a travel buddy who knows how to lift the mood
|
103 |
SAMPLE DIALOGUES
|
104 |
When a user opens the chatbot for the first time:
|
105 |
User: Hello?
|
106 |
+
Pixity: Hi hi 👋 I've been waiting for you! Ready to explore Da Nang together? I've got tips, tricks, and a tiny bit of magic 🎒✨
|
107 |
|
108 |
Return Format:
|
109 |
Respond in friendly, natural, concise and use only English like a real tour guide.
|
|
|
364 |
"services": services,
|
365 |
"retrieval_config": retrieval_config,
|
366 |
"timestamp": datetime.now().isoformat()
|
367 |
+
}
|
368 |
+
|
369 |
+
# Chat Engine endpoints
|
370 |
+
@router.get("/chat-engine", response_model=List[ChatEngineResponse], tags=["Chat Engine"])
|
371 |
+
async def get_chat_engines(
|
372 |
+
skip: int = 0,
|
373 |
+
limit: int = 100,
|
374 |
+
status: Optional[str] = None,
|
375 |
+
db: Session = Depends(get_db)
|
376 |
+
):
|
377 |
+
"""
|
378 |
+
Lấy danh sách tất cả chat engines.
|
379 |
+
|
380 |
+
- **skip**: Số lượng items bỏ qua
|
381 |
+
- **limit**: Số lượng items tối đa trả về
|
382 |
+
- **status**: Lọc theo trạng thái (ví dụ: 'active', 'inactive')
|
383 |
+
"""
|
384 |
+
try:
|
385 |
+
query = db.query(ChatEngine)
|
386 |
+
|
387 |
+
if status:
|
388 |
+
query = query.filter(ChatEngine.status == status)
|
389 |
+
|
390 |
+
engines = query.offset(skip).limit(limit).all()
|
391 |
+
return [ChatEngineResponse.model_validate(engine, from_attributes=True) for engine in engines]
|
392 |
+
except SQLAlchemyError as e:
|
393 |
+
logger.error(f"Database error retrieving chat engines: {e}")
|
394 |
+
raise HTTPException(status_code=500, detail=f"Lỗi database: {str(e)}")
|
395 |
+
except Exception as e:
|
396 |
+
logger.error(f"Error retrieving chat engines: {e}")
|
397 |
+
logger.error(traceback.format_exc())
|
398 |
+
raise HTTPException(status_code=500, detail=f"Lỗi khi lấy danh sách chat engines: {str(e)}")
|
399 |
+
|
400 |
+
@router.post("/chat-engine", response_model=ChatEngineResponse, status_code=status.HTTP_201_CREATED, tags=["Chat Engine"])
|
401 |
+
async def create_chat_engine(
|
402 |
+
engine: ChatEngineCreate,
|
403 |
+
db: Session = Depends(get_db)
|
404 |
+
):
|
405 |
+
"""
|
406 |
+
Tạo mới một chat engine.
|
407 |
+
|
408 |
+
- **name**: Tên của chat engine
|
409 |
+
- **answer_model**: Model được dùng để trả lời
|
410 |
+
- **system_prompt**: Prompt của hệ thống (optional)
|
411 |
+
- **empty_response**: Đoạn response khi không có thông tin (optional)
|
412 |
+
- **characteristic**: Tính cách của model (optional)
|
413 |
+
- **historical_sessions_number**: Số lượng các cặp tin nhắn trong history (default: 3)
|
414 |
+
- **use_public_information**: Cho phép sử dụng kiến thức bên ngoài (default: false)
|
415 |
+
- **similarity_top_k**: Số lượng documents tương tự (default: 3)
|
416 |
+
- **vector_distance_threshold**: Ngưỡng độ tương tự (default: 0.75)
|
417 |
+
- **grounding_threshold**: Ngưỡng grounding (default: 0.2)
|
418 |
+
- **pinecone_index_name**: Tên của vector database sử dụng (default: "testbot768")
|
419 |
+
- **status**: Trạng thái (default: "active")
|
420 |
+
"""
|
421 |
+
try:
|
422 |
+
# Create chat engine
|
423 |
+
db_engine = ChatEngine(**engine.model_dump())
|
424 |
+
|
425 |
+
db.add(db_engine)
|
426 |
+
db.commit()
|
427 |
+
db.refresh(db_engine)
|
428 |
+
|
429 |
+
return ChatEngineResponse.model_validate(db_engine, from_attributes=True)
|
430 |
+
except SQLAlchemyError as e:
|
431 |
+
db.rollback()
|
432 |
+
logger.error(f"Database error creating chat engine: {e}")
|
433 |
+
raise HTTPException(status_code=500, detail=f"Lỗi database: {str(e)}")
|
434 |
+
except Exception as e:
|
435 |
+
db.rollback()
|
436 |
+
logger.error(f"Error creating chat engine: {e}")
|
437 |
+
logger.error(traceback.format_exc())
|
438 |
+
raise HTTPException(status_code=500, detail=f"Lỗi khi tạo chat engine: {str(e)}")
|
439 |
+
|
440 |
+
@router.get("/chat-engine/{engine_id}", response_model=ChatEngineResponse, tags=["Chat Engine"])
|
441 |
+
async def get_chat_engine(
|
442 |
+
engine_id: int = Path(..., gt=0, description="ID của chat engine"),
|
443 |
+
db: Session = Depends(get_db)
|
444 |
+
):
|
445 |
+
"""
|
446 |
+
Lấy thông tin chi tiết của một chat engine theo ID.
|
447 |
+
|
448 |
+
- **engine_id**: ID của chat engine
|
449 |
+
"""
|
450 |
+
try:
|
451 |
+
engine = db.query(ChatEngine).filter(ChatEngine.id == engine_id).first()
|
452 |
+
if not engine:
|
453 |
+
raise HTTPException(status_code=404, detail=f"Không tìm thấy chat engine với ID {engine_id}")
|
454 |
+
|
455 |
+
return ChatEngineResponse.model_validate(engine, from_attributes=True)
|
456 |
+
except HTTPException:
|
457 |
+
raise
|
458 |
+
except Exception as e:
|
459 |
+
logger.error(f"Error retrieving chat engine: {e}")
|
460 |
+
logger.error(traceback.format_exc())
|
461 |
+
raise HTTPException(status_code=500, detail=f"Lỗi khi lấy thông tin chat engine: {str(e)}")
|
462 |
+
|
463 |
+
@router.put("/chat-engine/{engine_id}", response_model=ChatEngineResponse, tags=["Chat Engine"])
|
464 |
+
async def update_chat_engine(
|
465 |
+
engine_id: int = Path(..., gt=0, description="ID của chat engine"),
|
466 |
+
engine_update: ChatEngineUpdate = Body(...),
|
467 |
+
db: Session = Depends(get_db)
|
468 |
+
):
|
469 |
+
"""
|
470 |
+
Cập nhật thông tin của một chat engine.
|
471 |
+
|
472 |
+
- **engine_id**: ID của chat engine
|
473 |
+
- **engine_update**: Dữ liệu cập nhật
|
474 |
+
"""
|
475 |
+
try:
|
476 |
+
db_engine = db.query(ChatEngine).filter(ChatEngine.id == engine_id).first()
|
477 |
+
if not db_engine:
|
478 |
+
raise HTTPException(status_code=404, detail=f"Không tìm thấy chat engine với ID {engine_id}")
|
479 |
+
|
480 |
+
# Update fields if provided
|
481 |
+
update_data = engine_update.model_dump(exclude_unset=True)
|
482 |
+
for key, value in update_data.items():
|
483 |
+
if value is not None:
|
484 |
+
setattr(db_engine, key, value)
|
485 |
+
|
486 |
+
# Update last_modified timestamp
|
487 |
+
db_engine.last_modified = datetime.utcnow()
|
488 |
+
|
489 |
+
db.commit()
|
490 |
+
db.refresh(db_engine)
|
491 |
+
|
492 |
+
return ChatEngineResponse.model_validate(db_engine, from_attributes=True)
|
493 |
+
except HTTPException:
|
494 |
+
raise
|
495 |
+
except SQLAlchemyError as e:
|
496 |
+
db.rollback()
|
497 |
+
logger.error(f"Database error updating chat engine: {e}")
|
498 |
+
raise HTTPException(status_code=500, detail=f"Lỗi database: {str(e)}")
|
499 |
+
except Exception as e:
|
500 |
+
db.rollback()
|
501 |
+
logger.error(f"Error updating chat engine: {e}")
|
502 |
+
logger.error(traceback.format_exc())
|
503 |
+
raise HTTPException(status_code=500, detail=f"Lỗi khi cập nhật chat engine: {str(e)}")
|
504 |
+
|
505 |
+
@router.delete("/chat-engine/{engine_id}", response_model=dict, tags=["Chat Engine"])
|
506 |
+
async def delete_chat_engine(
|
507 |
+
engine_id: int = Path(..., gt=0, description="ID của chat engine"),
|
508 |
+
db: Session = Depends(get_db)
|
509 |
+
):
|
510 |
+
"""
|
511 |
+
Xóa một chat engine.
|
512 |
+
|
513 |
+
- **engine_id**: ID của chat engine
|
514 |
+
"""
|
515 |
+
try:
|
516 |
+
db_engine = db.query(ChatEngine).filter(ChatEngine.id == engine_id).first()
|
517 |
+
if not db_engine:
|
518 |
+
raise HTTPException(status_code=404, detail=f"Không tìm thấy chat engine với ID {engine_id}")
|
519 |
+
|
520 |
+
# Delete engine
|
521 |
+
db.delete(db_engine)
|
522 |
+
db.commit()
|
523 |
+
|
524 |
+
return {"message": f"Chat engine với ID {engine_id} đã được xóa thành công"}
|
525 |
+
except HTTPException:
|
526 |
+
raise
|
527 |
+
except SQLAlchemyError as e:
|
528 |
+
db.rollback()
|
529 |
+
logger.error(f"Database error deleting chat engine: {e}")
|
530 |
+
raise HTTPException(status_code=500, detail=f"Lỗi database: {str(e)}")
|
531 |
+
except Exception as e:
|
532 |
+
db.rollback()
|
533 |
+
logger.error(f"Error deleting chat engine: {e}")
|
534 |
+
logger.error(traceback.format_exc())
|
535 |
+
raise HTTPException(status_code=500, detail=f"Lỗi khi xóa chat engine: {str(e)}")
|
536 |
+
|
537 |
+
@timer_decorator
|
538 |
+
@router.post("/chat-with-engine/{engine_id}", response_model=ChatResponse, tags=["Chat Engine"])
|
539 |
+
async def chat_with_engine(
|
540 |
+
engine_id: int = Path(..., gt=0, description="ID của chat engine"),
|
541 |
+
request: ChatWithEngineRequest = Body(...),
|
542 |
+
background_tasks: BackgroundTasks = None,
|
543 |
+
db: Session = Depends(get_db)
|
544 |
+
):
|
545 |
+
"""
|
546 |
+
Tương tác với một chat engine cụ thể.
|
547 |
+
|
548 |
+
- **engine_id**: ID của chat engine
|
549 |
+
- **user_id**: ID của người dùng
|
550 |
+
- **question**: Câu hỏi của người dùng
|
551 |
+
- **include_history**: Có sử dụng lịch sử chat hay không
|
552 |
+
- **session_id**: ID session (optional)
|
553 |
+
- **first_name**: Tên của người dùng (optional)
|
554 |
+
- **last_name**: Họ của người dùng (optional)
|
555 |
+
- **username**: Username của người dùng (optional)
|
556 |
+
"""
|
557 |
+
start_time = time.time()
|
558 |
+
try:
|
559 |
+
# Lấy cache
|
560 |
+
cache = get_cache()
|
561 |
+
cache_key = get_chat_engine_cache_key(engine_id)
|
562 |
+
|
563 |
+
# Kiểm tra cache trước
|
564 |
+
engine = cache.get(cache_key)
|
565 |
+
if not engine:
|
566 |
+
logger.debug(f"Cache miss for engine ID {engine_id}, fetching from database")
|
567 |
+
# Nếu không có trong cache, truy vấn database
|
568 |
+
engine = db.query(ChatEngine).filter(ChatEngine.id == engine_id).first()
|
569 |
+
if not engine:
|
570 |
+
raise HTTPException(status_code=404, detail=f"Không tìm thấy chat engine với ID {engine_id}")
|
571 |
+
|
572 |
+
# Lưu vào cache
|
573 |
+
cache.set(cache_key, engine, CHAT_ENGINE_CACHE_TTL)
|
574 |
+
else:
|
575 |
+
logger.debug(f"Cache hit for engine ID {engine_id}")
|
576 |
+
|
577 |
+
# Kiểm tra trạng thái của engine
|
578 |
+
if engine.status != "active":
|
579 |
+
raise HTTPException(status_code=400, detail=f"Chat engine với ID {engine_id} không hoạt động")
|
580 |
+
|
581 |
+
# Lưu tin nhắn người dùng
|
582 |
+
session_id = request.session_id or f"{request.user_id}_{datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}"
|
583 |
+
|
584 |
+
# Cache các tham số cấu hình retriever
|
585 |
+
retriever_cache_key = get_retriever_cache_key(engine_id)
|
586 |
+
retriever_params = cache.get(retriever_cache_key)
|
587 |
+
|
588 |
+
if not retriever_params:
|
589 |
+
# Nếu không có trong cache, tạo mới và lưu cache
|
590 |
+
retriever_params = {
|
591 |
+
"index_name": engine.pinecone_index_name,
|
592 |
+
"top_k": engine.similarity_top_k,
|
593 |
+
"limit_k": engine.similarity_top_k * 2, # Mặc định lấy gấp đôi top_k
|
594 |
+
"similarity_metric": DEFAULT_SIMILARITY_METRIC,
|
595 |
+
"similarity_threshold": engine.vector_distance_threshold
|
596 |
+
}
|
597 |
+
cache.set(retriever_cache_key, retriever_params, RETRIEVER_CACHE_TTL)
|
598 |
+
|
599 |
+
# Khởi tạo retriever với các tham số từ cache
|
600 |
+
retriever = get_chain(**retriever_params)
|
601 |
+
if not retriever:
|
602 |
+
raise HTTPException(status_code=500, detail="Không thể khởi tạo retriever")
|
603 |
+
|
604 |
+
# Lấy lịch sử chat nếu cần
|
605 |
+
chat_history = ""
|
606 |
+
if request.include_history and engine.historical_sessions_number > 0:
|
607 |
+
chat_history = get_chat_history(request.user_id, n=engine.historical_sessions_number)
|
608 |
+
logger.info(f"Sử dụng lịch sử chat: {chat_history[:100]}...")
|
609 |
+
|
610 |
+
# Cache các tham số cấu hình model
|
611 |
+
model_cache_key = get_model_config_cache_key(engine.answer_model)
|
612 |
+
model_config = cache.get(model_cache_key)
|
613 |
+
|
614 |
+
if not model_config:
|
615 |
+
# Nếu không có trong cache, tạo mới và lưu cache
|
616 |
+
generation_config = {
|
617 |
+
"temperature": 0.9,
|
618 |
+
"top_p": 1,
|
619 |
+
"top_k": 1,
|
620 |
+
"max_output_tokens": 2048,
|
621 |
+
}
|
622 |
+
|
623 |
+
safety_settings = [
|
624 |
+
{
|
625 |
+
"category": "HARM_CATEGORY_HARASSMENT",
|
626 |
+
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
|
627 |
+
},
|
628 |
+
{
|
629 |
+
"category": "HARM_CATEGORY_HATE_SPEECH",
|
630 |
+
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
|
631 |
+
},
|
632 |
+
{
|
633 |
+
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
634 |
+
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
|
635 |
+
},
|
636 |
+
{
|
637 |
+
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
|
638 |
+
"threshold": "BLOCK_MEDIUM_AND_ABOVE"
|
639 |
+
},
|
640 |
+
]
|
641 |
+
|
642 |
+
model_config = {
|
643 |
+
"model_name": engine.answer_model,
|
644 |
+
"generation_config": generation_config,
|
645 |
+
"safety_settings": safety_settings
|
646 |
+
}
|
647 |
+
|
648 |
+
cache.set(model_cache_key, model_config, MODEL_CONFIG_CACHE_TTL)
|
649 |
+
|
650 |
+
# Khởi tạo Gemini model từ cấu hình đã cache
|
651 |
+
model = genai.GenerativeModel(**model_config)
|
652 |
+
|
653 |
+
# Sử dụng fix_request để tinh chỉnh câu hỏi
|
654 |
+
prompt_request = fix_request.format(
|
655 |
+
question=request.question,
|
656 |
+
chat_history=chat_history
|
657 |
+
)
|
658 |
+
|
659 |
+
# Log thời gian bắt đầu final_request
|
660 |
+
final_request_start_time = time.time()
|
661 |
+
final_request = model.generate_content(prompt_request)
|
662 |
+
# Log thời gian hoàn thành final_request
|
663 |
+
logger.info(f"Fixed Request: {final_request.text}")
|
664 |
+
logger.info(f"Thời gian sinh fixed request: {time.time() - final_request_start_time:.2f} giây")
|
665 |
+
|
666 |
+
# Lấy context từ retriever
|
667 |
+
retrieved_docs = retriever.invoke(final_request.text)
|
668 |
+
logger.info(f"Số lượng tài liệu lấy được: {len(retrieved_docs)}")
|
669 |
+
context = "\n".join([doc.page_content for doc in retrieved_docs])
|
670 |
+
|
671 |
+
# Tạo danh sách nguồn
|
672 |
+
sources = []
|
673 |
+
for doc in retrieved_docs:
|
674 |
+
source = None
|
675 |
+
metadata = {}
|
676 |
+
|
677 |
+
if hasattr(doc, 'metadata'):
|
678 |
+
source = doc.metadata.get('source', None)
|
679 |
+
# Extract score information
|
680 |
+
score = doc.metadata.get('score', None)
|
681 |
+
normalized_score = doc.metadata.get('normalized_score', None)
|
682 |
+
# Remove score info from metadata to avoid duplication
|
683 |
+
metadata = {k: v for k, v in doc.metadata.items()
|
684 |
+
if k not in ['text', 'source', 'score', 'normalized_score']}
|
685 |
+
|
686 |
+
sources.append(SourceDocument(
|
687 |
+
text=doc.page_content,
|
688 |
+
source=source,
|
689 |
+
score=score,
|
690 |
+
normalized_score=normalized_score,
|
691 |
+
metadata=metadata
|
692 |
+
))
|
693 |
+
|
694 |
+
# Cache prompt template parameters
|
695 |
+
prompt_template_cache_key = get_prompt_template_cache_key(engine_id)
|
696 |
+
prompt_template_params = cache.get(prompt_template_cache_key)
|
697 |
+
|
698 |
+
if not prompt_template_params:
|
699 |
+
# Tạo prompt động dựa trên thông tin chat engine
|
700 |
+
system_prompt_part = engine.system_prompt or ""
|
701 |
+
empty_response_part = engine.empty_response or "I'm sorry. I don't have information about that."
|
702 |
+
characteristic_part = engine.characteristic or ""
|
703 |
+
use_public_info_part = "You can use your own knowledge." if engine.use_public_information else "Only use the information provided in the context to answer. If you do not have enough information, respond with the empty response."
|
704 |
+
|
705 |
+
prompt_template_params = {
|
706 |
+
"system_prompt_part": system_prompt_part,
|
707 |
+
"empty_response_part": empty_response_part,
|
708 |
+
"characteristic_part": characteristic_part,
|
709 |
+
"use_public_info_part": use_public_info_part
|
710 |
+
}
|
711 |
+
|
712 |
+
cache.set(prompt_template_cache_key, prompt_template_params, PROMPT_TEMPLATE_CACHE_TTL)
|
713 |
+
|
714 |
+
# Tạo final_prompt từ cache
|
715 |
+
final_prompt = f"""
|
716 |
+
{prompt_template_params['system_prompt_part']}
|
717 |
+
|
718 |
+
Your characteristics:
|
719 |
+
{prompt_template_params['characteristic_part']}
|
720 |
+
|
721 |
+
When you don't have enough information:
|
722 |
+
{prompt_template_params['empty_response_part']}
|
723 |
+
|
724 |
+
Knowledge usage instructions:
|
725 |
+
{prompt_template_params['use_public_info_part']}
|
726 |
+
|
727 |
+
Context:
|
728 |
+
{context}
|
729 |
+
|
730 |
+
Conversation History:
|
731 |
+
{chat_history}
|
732 |
+
|
733 |
+
User message:
|
734 |
+
{request.question}
|
735 |
+
|
736 |
+
Your response:
|
737 |
+
"""
|
738 |
+
|
739 |
+
logger.info(f"Final prompt: {final_prompt}")
|
740 |
+
|
741 |
+
# Sinh câu trả lời
|
742 |
+
response = model.generate_content(final_prompt)
|
743 |
+
answer = response.text
|
744 |
+
|
745 |
+
# Tính thời gian xử lý
|
746 |
+
processing_time = time.time() - start_time
|
747 |
+
|
748 |
+
# Tạo response object
|
749 |
+
chat_response = ChatResponse(
|
750 |
+
answer=answer,
|
751 |
+
processing_time=processing_time
|
752 |
+
)
|
753 |
+
|
754 |
+
# Trả về response
|
755 |
+
return chat_response
|
756 |
+
except Exception as e:
|
757 |
+
logger.error(f"Lỗi khi xử lý chat request: {e}")
|
758 |
+
logger.error(traceback.format_exc())
|
759 |
+
raise HTTPException(status_code=500, detail=f"Lỗi khi xử lý chat request: {str(e)}")
|
760 |
+
|
761 |
+
@router.get("/cache/stats", tags=["Cache"])
|
762 |
+
async def get_cache_stats():
|
763 |
+
"""
|
764 |
+
Lấy thống kê về cache.
|
765 |
+
|
766 |
+
Trả về thông tin về số lượng item trong cache, bộ nhớ sử dụng, v.v.
|
767 |
+
"""
|
768 |
+
try:
|
769 |
+
cache = get_cache()
|
770 |
+
stats = cache.stats()
|
771 |
+
|
772 |
+
# Bổ sung thông tin về cấu hình
|
773 |
+
stats.update({
|
774 |
+
"chat_engine_ttl": CHAT_ENGINE_CACHE_TTL,
|
775 |
+
"model_config_ttl": MODEL_CONFIG_CACHE_TTL,
|
776 |
+
"retriever_ttl": RETRIEVER_CACHE_TTL,
|
777 |
+
"prompt_template_ttl": PROMPT_TEMPLATE_CACHE_TTL
|
778 |
+
})
|
779 |
+
|
780 |
+
return stats
|
781 |
+
except Exception as e:
|
782 |
+
logger.error(f"Lỗi khi lấy thống kê cache: {e}")
|
783 |
+
logger.error(traceback.format_exc())
|
784 |
+
raise HTTPException(status_code=500, detail=f"Lỗi khi lấy thống kê cache: {str(e)}")
|
785 |
+
|
786 |
+
@router.delete("/cache", tags=["Cache"])
|
787 |
+
async def clear_cache(key: Optional[str] = None):
|
788 |
+
"""
|
789 |
+
Xóa cache.
|
790 |
+
|
791 |
+
- **key**: Key cụ thể cần xóa. Nếu không có, xóa toàn bộ cache.
|
792 |
+
"""
|
793 |
+
try:
|
794 |
+
cache = get_cache()
|
795 |
+
|
796 |
+
if key:
|
797 |
+
# Xóa một key cụ thể
|
798 |
+
success = cache.delete(key)
|
799 |
+
if success:
|
800 |
+
return {"message": f"Đã xóa cache cho key: {key}"}
|
801 |
+
else:
|
802 |
+
return {"message": f"Không tìm thấy key: {key} trong cache"}
|
803 |
+
else:
|
804 |
+
# Xóa toàn bộ cache
|
805 |
+
cache.clear()
|
806 |
+
return {"message": "Đã xóa toàn bộ cache"}
|
807 |
+
except Exception as e:
|
808 |
+
logger.error(f"Lỗi khi xóa cache: {e}")
|
809 |
+
logger.error(traceback.format_exc())
|
810 |
+
raise HTTPException(status_code=500, detail=f"Lỗi khi xóa cache: {str(e)}")
|
app/database/models.py
CHANGED
@@ -125,7 +125,6 @@ class VectorStatus(Base):
|
|
125 |
document_id = Column(Integer, ForeignKey("document.id"), nullable=False)
|
126 |
vector_database_id = Column(Integer, ForeignKey("vector_database.id"), nullable=False)
|
127 |
vector_id = Column(String, nullable=True)
|
128 |
-
document_name = Column(String, nullable=True)
|
129 |
status = Column(String, default="pending")
|
130 |
error_message = Column(String, nullable=True)
|
131 |
embedded_at = Column(DateTime, nullable=True)
|
@@ -156,10 +155,13 @@ class ChatEngine(Base):
|
|
156 |
answer_model = Column(String, nullable=False)
|
157 |
system_prompt = Column(Text, nullable=True)
|
158 |
empty_response = Column(String, nullable=True)
|
|
|
|
|
159 |
similarity_top_k = Column(Integer, default=3)
|
160 |
vector_distance_threshold = Column(Float, default=0.75)
|
161 |
grounding_threshold = Column(Float, default=0.2)
|
162 |
use_public_information = Column(Boolean, default=False)
|
|
|
163 |
status = Column(String, default="active")
|
164 |
created_at = Column(DateTime, server_default=func.now())
|
165 |
last_modified = Column(DateTime, server_default=func.now(), onupdate=func.now())
|
|
|
125 |
document_id = Column(Integer, ForeignKey("document.id"), nullable=False)
|
126 |
vector_database_id = Column(Integer, ForeignKey("vector_database.id"), nullable=False)
|
127 |
vector_id = Column(String, nullable=True)
|
|
|
128 |
status = Column(String, default="pending")
|
129 |
error_message = Column(String, nullable=True)
|
130 |
embedded_at = Column(DateTime, nullable=True)
|
|
|
155 |
answer_model = Column(String, nullable=False)
|
156 |
system_prompt = Column(Text, nullable=True)
|
157 |
empty_response = Column(String, nullable=True)
|
158 |
+
characteristic = Column(Text, nullable=True)
|
159 |
+
historical_sessions_number = Column(Integer, default=3)
|
160 |
similarity_top_k = Column(Integer, default=3)
|
161 |
vector_distance_threshold = Column(Float, default=0.75)
|
162 |
grounding_threshold = Column(Float, default=0.2)
|
163 |
use_public_information = Column(Boolean, default=False)
|
164 |
+
pinecone_index_name = Column(String, default="testbot768")
|
165 |
status = Column(String, default="active")
|
166 |
created_at = Column(DateTime, server_default=func.now())
|
167 |
last_modified = Column(DateTime, server_default=func.now(), onupdate=func.now())
|
app/models/rag_models.py
CHANGED
@@ -1,5 +1,7 @@
|
|
1 |
from pydantic import BaseModel, Field
|
2 |
from typing import Optional, List, Dict, Any
|
|
|
|
|
3 |
|
4 |
class ChatRequest(BaseModel):
|
5 |
"""Request model for chat endpoint"""
|
@@ -12,7 +14,7 @@ class ChatRequest(BaseModel):
|
|
12 |
similarity_top_k: int = Field(6, description="Number of top similar documents to return (after filtering)")
|
13 |
limit_k: int = Field(10, description="Maximum number of documents to retrieve from vector store")
|
14 |
similarity_metric: str = Field("cosine", description="Similarity metric to use (cosine, dotproduct, euclidean)")
|
15 |
-
similarity_threshold: float = Field(0.
|
16 |
|
17 |
# User information
|
18 |
session_id: Optional[str] = Field(None, description="Session ID for tracking conversations")
|
@@ -65,4 +67,58 @@ class UserMessageModel(BaseModel):
|
|
65 |
similarity_top_k: Optional[int] = Field(None, description="Number of top similar documents to return (after filtering)")
|
66 |
limit_k: Optional[int] = Field(None, description="Maximum number of documents to retrieve from vector store")
|
67 |
similarity_metric: Optional[str] = Field(None, description="Similarity metric to use (cosine, dotproduct, euclidean)")
|
68 |
-
similarity_threshold: Optional[float] = Field(None, description="Threshold for vector similarity (0-1)")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
from pydantic import BaseModel, Field
|
2 |
from typing import Optional, List, Dict, Any
|
3 |
+
from datetime import datetime
|
4 |
+
from pydantic import ConfigDict
|
5 |
|
6 |
class ChatRequest(BaseModel):
|
7 |
"""Request model for chat endpoint"""
|
|
|
14 |
similarity_top_k: int = Field(6, description="Number of top similar documents to return (after filtering)")
|
15 |
limit_k: int = Field(10, description="Maximum number of documents to retrieve from vector store")
|
16 |
similarity_metric: str = Field("cosine", description="Similarity metric to use (cosine, dotproduct, euclidean)")
|
17 |
+
similarity_threshold: float = Field(0.0, description="Threshold for vector similarity (0-1)")
|
18 |
|
19 |
# User information
|
20 |
session_id: Optional[str] = Field(None, description="Session ID for tracking conversations")
|
|
|
67 |
similarity_top_k: Optional[int] = Field(None, description="Number of top similar documents to return (after filtering)")
|
68 |
limit_k: Optional[int] = Field(None, description="Maximum number of documents to retrieve from vector store")
|
69 |
similarity_metric: Optional[str] = Field(None, description="Similarity metric to use (cosine, dotproduct, euclidean)")
|
70 |
+
similarity_threshold: Optional[float] = Field(None, description="Threshold for vector similarity (0-1)")
|
71 |
+
|
72 |
+
class ChatEngineBase(BaseModel):
|
73 |
+
"""Base model cho chat engine"""
|
74 |
+
name: str = Field(..., description="Tên của chat engine")
|
75 |
+
answer_model: str = Field(..., description="Model được dùng để trả lời")
|
76 |
+
system_prompt: Optional[str] = Field(None, description="Prompt của hệ thống, được đưa vào phần đầu tiên của final_prompt")
|
77 |
+
empty_response: Optional[str] = Field(None, description="Đoạn response khi answer model không có thông tin về câu hỏi")
|
78 |
+
characteristic: Optional[str] = Field(None, description="Tính cách của model khi trả lời câu hỏi")
|
79 |
+
historical_sessions_number: int = Field(3, description="Số lượng các cặp tin nhắn trong history được đưa vào final prompt")
|
80 |
+
use_public_information: bool = Field(False, description="Yes nếu answer model được quyền trả về thông tin mà nó có")
|
81 |
+
similarity_top_k: int = Field(3, description="Số lượng top similar documents để trả về")
|
82 |
+
vector_distance_threshold: float = Field(0.75, description="Threshold cho vector similarity")
|
83 |
+
grounding_threshold: float = Field(0.2, description="Threshold cho grounding")
|
84 |
+
pinecone_index_name: str = Field("testbot768", description="Vector database mà model được quyền sử dụng")
|
85 |
+
status: str = Field("active", description="Trạng thái của chat engine")
|
86 |
+
|
87 |
+
class ChatEngineCreate(ChatEngineBase):
|
88 |
+
"""Model cho việc tạo chat engine mới"""
|
89 |
+
pass
|
90 |
+
|
91 |
+
class ChatEngineUpdate(BaseModel):
|
92 |
+
"""Model cho việc cập nhật chat engine"""
|
93 |
+
name: Optional[str] = None
|
94 |
+
answer_model: Optional[str] = None
|
95 |
+
system_prompt: Optional[str] = None
|
96 |
+
empty_response: Optional[str] = None
|
97 |
+
characteristic: Optional[str] = None
|
98 |
+
historical_sessions_number: Optional[int] = None
|
99 |
+
use_public_information: Optional[bool] = None
|
100 |
+
similarity_top_k: Optional[int] = None
|
101 |
+
vector_distance_threshold: Optional[float] = None
|
102 |
+
grounding_threshold: Optional[float] = None
|
103 |
+
pinecone_index_name: Optional[str] = None
|
104 |
+
status: Optional[str] = None
|
105 |
+
|
106 |
+
class ChatEngineResponse(ChatEngineBase):
|
107 |
+
"""Response model cho chat engine"""
|
108 |
+
id: int
|
109 |
+
created_at: datetime
|
110 |
+
last_modified: datetime
|
111 |
+
|
112 |
+
model_config = ConfigDict(from_attributes=True)
|
113 |
+
|
114 |
+
class ChatWithEngineRequest(BaseModel):
|
115 |
+
"""Request model cho endpoint chat-with-engine"""
|
116 |
+
user_id: str = Field(..., description="User ID from Telegram")
|
117 |
+
question: str = Field(..., description="User's question")
|
118 |
+
include_history: bool = Field(True, description="Whether to include user history in prompt")
|
119 |
+
|
120 |
+
# User information
|
121 |
+
session_id: Optional[str] = Field(None, description="Session ID for tracking conversations")
|
122 |
+
first_name: Optional[str] = Field(None, description="User's first name")
|
123 |
+
last_name: Optional[str] = Field(None, description="User's last name")
|
124 |
+
username: Optional[str] = Field(None, description="User's username")
|
app/utils/cache_config.py
ADDED
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""
|
2 |
+
Module cấu hình cho cache.
|
3 |
+
|
4 |
+
Module này chứa các tham số cấu hình và constants liên quan đến cache.
|
5 |
+
"""
|
6 |
+
|
7 |
+
import os
|
8 |
+
from dotenv import load_dotenv
|
9 |
+
|
10 |
+
# Load biến môi trường
|
11 |
+
load_dotenv()
|
12 |
+
|
13 |
+
# Cấu hình cache từ biến môi trường, có thể override bằng .env file
|
14 |
+
CACHE_TTL_SECONDS = int(os.getenv("CACHE_TTL_SECONDS", "300")) # Mặc định 5 phút
|
15 |
+
CACHE_CLEANUP_INTERVAL = int(os.getenv("CACHE_CLEANUP_INTERVAL", "60")) # Mặc định 1 phút
|
16 |
+
CACHE_MAX_SIZE = int(os.getenv("CACHE_MAX_SIZE", "1000")) # Mặc định 1000 phần tử
|
17 |
+
|
18 |
+
# Cấu hình cho loại cache cụ thể
|
19 |
+
CHAT_ENGINE_CACHE_TTL = int(os.getenv("CHAT_ENGINE_CACHE_TTL", str(CACHE_TTL_SECONDS)))
|
20 |
+
MODEL_CONFIG_CACHE_TTL = int(os.getenv("MODEL_CONFIG_CACHE_TTL", str(CACHE_TTL_SECONDS)))
|
21 |
+
RETRIEVER_CACHE_TTL = int(os.getenv("RETRIEVER_CACHE_TTL", str(CACHE_TTL_SECONDS)))
|
22 |
+
PROMPT_TEMPLATE_CACHE_TTL = int(os.getenv("PROMPT_TEMPLATE_CACHE_TTL", str(CACHE_TTL_SECONDS)))
|
23 |
+
|
24 |
+
# Cache keys prefix
|
25 |
+
CHAT_ENGINE_CACHE_PREFIX = "chat_engine:"
|
26 |
+
MODEL_CONFIG_CACHE_PREFIX = "model_config:"
|
27 |
+
RETRIEVER_CACHE_PREFIX = "retriever:"
|
28 |
+
PROMPT_TEMPLATE_CACHE_PREFIX = "prompt_template:"
|
29 |
+
|
30 |
+
# Hàm helper để tạo cache key
|
31 |
+
def get_chat_engine_cache_key(engine_id: int) -> str:
|
32 |
+
"""Tạo cache key cho chat engine"""
|
33 |
+
return f"{CHAT_ENGINE_CACHE_PREFIX}{engine_id}"
|
34 |
+
|
35 |
+
def get_model_config_cache_key(model_name: str) -> str:
|
36 |
+
"""Tạo cache key cho model config"""
|
37 |
+
return f"{MODEL_CONFIG_CACHE_PREFIX}{model_name}"
|
38 |
+
|
39 |
+
def get_retriever_cache_key(engine_id: int) -> str:
|
40 |
+
"""Tạo cache key cho retriever"""
|
41 |
+
return f"{RETRIEVER_CACHE_PREFIX}{engine_id}"
|
42 |
+
|
43 |
+
def get_prompt_template_cache_key(engine_id: int) -> str:
|
44 |
+
"""Tạo cache key cho prompt template"""
|
45 |
+
return f"{PROMPT_TEMPLATE_CACHE_PREFIX}{engine_id}"
|
beach_request.json
DELETED
Binary file (470 Bytes)
|
|
chat_request.json
DELETED
Binary file (472 Bytes)
|
|
pytest.ini
DELETED
@@ -1,12 +0,0 @@
|
|
1 |
-
[pytest]
|
2 |
-
# Bỏ qua cảnh báo về anyio module và các cảnh báo vận hành nội bộ
|
3 |
-
filterwarnings =
|
4 |
-
ignore::pytest.PytestAssertRewriteWarning:.*anyio
|
5 |
-
ignore:.*general_plain_validator_function.* is deprecated.*:DeprecationWarning
|
6 |
-
ignore:.*with_info_plain_validator_function.*:DeprecationWarning
|
7 |
-
|
8 |
-
# Cấu hình cơ bản khác
|
9 |
-
testpaths = tests
|
10 |
-
python_files = test_*.py
|
11 |
-
python_classes = Test*
|
12 |
-
python_functions = test_*
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
test_body.json
DELETED
Binary file (864 Bytes)
|
|
test_rag_api.py
DELETED
@@ -1,263 +0,0 @@
|
|
1 |
-
import requests
|
2 |
-
import json
|
3 |
-
import psycopg2
|
4 |
-
import os
|
5 |
-
from dotenv import load_dotenv
|
6 |
-
|
7 |
-
# Load .env file if it exists
|
8 |
-
load_dotenv()
|
9 |
-
|
10 |
-
# PostgreSQL connection parameters
|
11 |
-
# For testing purposes, let's use localhost PostgreSQL if not available from environment
|
12 |
-
DB_CONNECTION_MODE = os.getenv("DB_CONNECTION_MODE", "local")
|
13 |
-
DATABASE_URL = os.getenv("AIVEN_DB_URL")
|
14 |
-
|
15 |
-
# Default test parameters - will be used if env vars not set
|
16 |
-
DEFAULT_DB_USER = "postgres"
|
17 |
-
DEFAULT_DB_PASSWORD = "postgres"
|
18 |
-
DEFAULT_DB_HOST = "localhost"
|
19 |
-
DEFAULT_DB_PORT = "5432"
|
20 |
-
DEFAULT_DB_NAME = "pixity"
|
21 |
-
|
22 |
-
# Parse DATABASE_URL if available, otherwise use defaults
|
23 |
-
if DATABASE_URL:
|
24 |
-
try:
|
25 |
-
# Extract credentials and host info
|
26 |
-
credentials, rest = DATABASE_URL.split("@")
|
27 |
-
user_pass = credentials.split("://")[1]
|
28 |
-
host_port_db = rest.split("/")
|
29 |
-
|
30 |
-
# Split user/pass and host/port
|
31 |
-
if ":" in user_pass:
|
32 |
-
user, password = user_pass.split(":")
|
33 |
-
else:
|
34 |
-
user, password = user_pass, ""
|
35 |
-
|
36 |
-
host_port = host_port_db[0]
|
37 |
-
if ":" in host_port:
|
38 |
-
host, port = host_port.split(":")
|
39 |
-
else:
|
40 |
-
host, port = host_port, "5432"
|
41 |
-
|
42 |
-
# Get database name
|
43 |
-
dbname = host_port_db[1]
|
44 |
-
if "?" in dbname:
|
45 |
-
dbname = dbname.split("?")[0]
|
46 |
-
|
47 |
-
print(f"Parsed connection parameters: host={host}, port={port}, dbname={dbname}, user={user}")
|
48 |
-
except Exception as e:
|
49 |
-
print(f"Error parsing DATABASE_URL: {e}")
|
50 |
-
print("Using default connection parameters")
|
51 |
-
user = DEFAULT_DB_USER
|
52 |
-
password = DEFAULT_DB_PASSWORD
|
53 |
-
host = DEFAULT_DB_HOST
|
54 |
-
port = DEFAULT_DB_PORT
|
55 |
-
dbname = DEFAULT_DB_NAME
|
56 |
-
else:
|
57 |
-
print("No DATABASE_URL found. Using default connection parameters")
|
58 |
-
user = DEFAULT_DB_USER
|
59 |
-
password = DEFAULT_DB_PASSWORD
|
60 |
-
host = DEFAULT_DB_HOST
|
61 |
-
port = DEFAULT_DB_PORT
|
62 |
-
dbname = DEFAULT_DB_NAME
|
63 |
-
|
64 |
-
# Execute direct SQL to add the column
|
65 |
-
def add_required_columns():
|
66 |
-
try:
|
67 |
-
print(f"Connecting to PostgreSQL: {host}:{port} database={dbname} user={user}")
|
68 |
-
# Connect to PostgreSQL
|
69 |
-
conn = psycopg2.connect(
|
70 |
-
user=user,
|
71 |
-
password=password,
|
72 |
-
host=host,
|
73 |
-
port=port,
|
74 |
-
dbname=dbname
|
75 |
-
)
|
76 |
-
|
77 |
-
# Create a cursor
|
78 |
-
cursor = conn.cursor()
|
79 |
-
|
80 |
-
# 1. Check if pinecone_index_name column already exists
|
81 |
-
cursor.execute("""
|
82 |
-
SELECT column_name
|
83 |
-
FROM information_schema.columns
|
84 |
-
WHERE table_name='chat_engine' AND column_name='pinecone_index_name';
|
85 |
-
""")
|
86 |
-
|
87 |
-
column_exists = cursor.fetchone()
|
88 |
-
|
89 |
-
if not column_exists:
|
90 |
-
print("Column 'pinecone_index_name' does not exist. Adding it...")
|
91 |
-
# Add the pinecone_index_name column to the chat_engine table
|
92 |
-
cursor.execute("""
|
93 |
-
ALTER TABLE chat_engine
|
94 |
-
ADD COLUMN pinecone_index_name VARCHAR NULL;
|
95 |
-
""")
|
96 |
-
conn.commit()
|
97 |
-
print("Column 'pinecone_index_name' added successfully!")
|
98 |
-
else:
|
99 |
-
print("Column 'pinecone_index_name' already exists.")
|
100 |
-
|
101 |
-
# 2. Check if characteristic column already exists
|
102 |
-
cursor.execute("""
|
103 |
-
SELECT column_name
|
104 |
-
FROM information_schema.columns
|
105 |
-
WHERE table_name='chat_engine' AND column_name='characteristic';
|
106 |
-
""")
|
107 |
-
|
108 |
-
characteristic_exists = cursor.fetchone()
|
109 |
-
|
110 |
-
if not characteristic_exists:
|
111 |
-
print("Column 'characteristic' does not exist. Adding it...")
|
112 |
-
# Add the characteristic column to the chat_engine table
|
113 |
-
cursor.execute("""
|
114 |
-
ALTER TABLE chat_engine
|
115 |
-
ADD COLUMN characteristic TEXT NULL;
|
116 |
-
""")
|
117 |
-
conn.commit()
|
118 |
-
print("Column 'characteristic' added successfully!")
|
119 |
-
else:
|
120 |
-
print("Column 'characteristic' already exists.")
|
121 |
-
|
122 |
-
# Close cursor and connection
|
123 |
-
cursor.close()
|
124 |
-
conn.close()
|
125 |
-
return True
|
126 |
-
except Exception as e:
|
127 |
-
print(f"Error accessing PostgreSQL: {e}")
|
128 |
-
print("Please make sure PostgreSQL is running and accessible.")
|
129 |
-
return False
|
130 |
-
|
131 |
-
# Base URL
|
132 |
-
base_url = "http://localhost:7860"
|
133 |
-
|
134 |
-
def test_create_engine():
|
135 |
-
"""Test creating a new chat engine"""
|
136 |
-
url = f"{base_url}/rag/chat-engine"
|
137 |
-
data = {
|
138 |
-
"name": "Test Engine",
|
139 |
-
"answer_model": "models/gemini-2.0-flash",
|
140 |
-
"system_prompt": "You are an AI assistant that helps users find information about Da Nang.",
|
141 |
-
"empty_response": "I don't have information about this question.",
|
142 |
-
"use_public_information": True,
|
143 |
-
"similarity_top_k": 5,
|
144 |
-
"vector_distance_threshold": 0.7,
|
145 |
-
"grounding_threshold": 0.2,
|
146 |
-
"pinecone_index_name": "testbot768",
|
147 |
-
"characteristic": "You are friendly, helpful, and concise. You use a warm and conversational tone, and occasionally add emojis to seem more personable. You always try to be specific in your answers and provide examples when relevant.",
|
148 |
-
"status": "active"
|
149 |
-
}
|
150 |
-
|
151 |
-
response = requests.post(url, json=data)
|
152 |
-
print(f"Create Engine Response Status: {response.status_code}")
|
153 |
-
if response.status_code == 201 or response.status_code == 200:
|
154 |
-
print(f"Successfully created engine: {response.json()}")
|
155 |
-
return response.json().get("id")
|
156 |
-
else:
|
157 |
-
print(f"Failed to create engine: {response.text}")
|
158 |
-
return None
|
159 |
-
|
160 |
-
def test_get_engine(engine_id):
|
161 |
-
"""Test getting a specific chat engine"""
|
162 |
-
url = f"{base_url}/rag/chat-engine/{engine_id}"
|
163 |
-
response = requests.get(url)
|
164 |
-
print(f"Get Engine Response Status: {response.status_code}")
|
165 |
-
if response.status_code == 200:
|
166 |
-
print(f"Engine details: {response.json()}")
|
167 |
-
else:
|
168 |
-
print(f"Failed to get engine: {response.text}")
|
169 |
-
|
170 |
-
def test_list_engines():
|
171 |
-
"""Test listing all chat engines"""
|
172 |
-
url = f"{base_url}/rag/chat-engines"
|
173 |
-
response = requests.get(url)
|
174 |
-
print(f"List Engines Response Status: {response.status_code}")
|
175 |
-
if response.status_code == 200:
|
176 |
-
engines = response.json()
|
177 |
-
print(f"Found {len(engines)} engines")
|
178 |
-
for engine in engines:
|
179 |
-
print(f" - ID: {engine.get('id')}, Name: {engine.get('name')}")
|
180 |
-
else:
|
181 |
-
print(f"Failed to list engines: {response.text}")
|
182 |
-
|
183 |
-
def test_update_engine(engine_id):
|
184 |
-
"""Test updating a chat engine"""
|
185 |
-
url = f"{base_url}/rag/chat-engine/{engine_id}"
|
186 |
-
data = {
|
187 |
-
"name": "Updated Test Engine",
|
188 |
-
"system_prompt": "You are an updated AI assistant for Da Nang information.",
|
189 |
-
"characteristic": "You speak in a very professional and formal tone. You are direct and to the point, avoiding unnecessary chatter. You prefer to use precise language and avoid colloquialisms."
|
190 |
-
}
|
191 |
-
|
192 |
-
response = requests.put(url, json=data)
|
193 |
-
print(f"Update Engine Response Status: {response.status_code}")
|
194 |
-
if response.status_code == 200:
|
195 |
-
print(f"Successfully updated engine: {response.json()}")
|
196 |
-
else:
|
197 |
-
print(f"Failed to update engine: {response.text}")
|
198 |
-
|
199 |
-
def test_chat_with_engine(engine_id):
|
200 |
-
"""Test chatting with a specific engine"""
|
201 |
-
url = f"{base_url}/rag/chat/{engine_id}"
|
202 |
-
data = {
|
203 |
-
"user_id": "test_user_123",
|
204 |
-
"question": "What are some popular attractions in Da Nang?",
|
205 |
-
"include_history": True,
|
206 |
-
"limit_k": 10,
|
207 |
-
"similarity_metric": "cosine",
|
208 |
-
"session_id": "test_session_123",
|
209 |
-
"first_name": "Test",
|
210 |
-
"last_name": "User",
|
211 |
-
"username": "testuser"
|
212 |
-
}
|
213 |
-
|
214 |
-
response = requests.post(url, json=data)
|
215 |
-
print(f"Chat With Engine Response Status: {response.status_code}")
|
216 |
-
if response.status_code == 200:
|
217 |
-
print(f"Chat response: {response.json()}")
|
218 |
-
else:
|
219 |
-
print(f"Failed to chat with engine: {response.text}")
|
220 |
-
|
221 |
-
def test_delete_engine(engine_id):
|
222 |
-
"""Test deleting a chat engine"""
|
223 |
-
url = f"{base_url}/rag/chat-engine/{engine_id}"
|
224 |
-
response = requests.delete(url)
|
225 |
-
print(f"Delete Engine Response Status: {response.status_code}")
|
226 |
-
if response.status_code == 204:
|
227 |
-
print(f"Successfully deleted engine with ID: {engine_id}")
|
228 |
-
else:
|
229 |
-
print(f"Failed to delete engine: {response.text}")
|
230 |
-
|
231 |
-
# Execute tests
|
232 |
-
if __name__ == "__main__":
|
233 |
-
print("First, let's add the missing columns to the database")
|
234 |
-
if add_required_columns():
|
235 |
-
print("\nStarting RAG Chat Engine API Tests")
|
236 |
-
print("---------------------------------")
|
237 |
-
|
238 |
-
# 1. Create a new engine
|
239 |
-
print("\n1. Testing Create Engine API")
|
240 |
-
engine_id = test_create_engine()
|
241 |
-
|
242 |
-
if engine_id:
|
243 |
-
# 2. Get engine details
|
244 |
-
print("\n2. Testing Get Engine API")
|
245 |
-
test_get_engine(engine_id)
|
246 |
-
|
247 |
-
# 3. List all engines
|
248 |
-
print("\n3. Testing List Engines API")
|
249 |
-
test_list_engines()
|
250 |
-
|
251 |
-
# 4. Update the engine
|
252 |
-
print("\n4. Testing Update Engine API")
|
253 |
-
test_update_engine(engine_id)
|
254 |
-
|
255 |
-
# 5. Chat with the engine
|
256 |
-
print("\n5. Testing Chat With Engine API")
|
257 |
-
test_chat_with_engine(engine_id)
|
258 |
-
|
259 |
-
# 6. Delete the engine
|
260 |
-
print("\n6. Testing Delete Engine API")
|
261 |
-
test_delete_engine(engine_id)
|
262 |
-
|
263 |
-
print("\nAPI Tests Completed")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
update_body.json
DELETED
Binary file (422 Bytes)
|
|