You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
117 lines
4.9 KiB
Python
117 lines
4.9 KiB
Python
import os
|
|
import json
|
|
import asyncio
|
|
import sys
|
|
import uuid
|
|
from datetime import datetime
|
|
from unittest.mock import MagicMock, AsyncMock
|
|
|
|
# Set PYTHONPATH to include the current directory
|
|
sys.path.append(os.getcwd())
|
|
|
|
from src.ircfa_interaction.schema import ChatRequest
|
|
from src.auth.model import UserBase
|
|
from src.ircfa_interaction.stream2 import stream_create
|
|
from src.ircfa_interaction.model import ChatHistory as ChatHistoryModel, Conversation as ConversationModel
|
|
from src.database.core import async_session, collector_async_session
|
|
|
|
# Token provided by USER
|
|
TOKEN = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJmcmVzaCI6ZmFsc2UsImlhdCI6MTc3NTEwNzEzOS4wODk5NjksImp0aSI6InRjOU9MbndhMkxFckFPaXF1YVk4aXciLCJ0eXBlIjoiYWNjZXNzIiwic3ViIjoiMTUzZjdlMmUtZjBlYS00NDBjLTk5Y2QtN2E3ODJhZThhNjMxIiwibmJmIjoxNzc1MTA3MTM5LCJjc3JmIjoiMmRiNzI3MWEtYTU5Mi00ODU4LWJjMzAtNDY4NGU2ODI1YWNjIiwiZXhwIjoyMDkwNDY3MTM5LCJub25jZSI6ImQzYWNkMDE5YjNiNzkyNWU3ZTY0MDFiMGQyMzhmNDQxIiwidXNlcm5hbWUiOiJkZXZhZG1pbiJ9.YjmpfEFuOAvdOyp9wpPAaykdKGkHFwbNayyVPNDrw_-13-noQwY76lbns-aGF3DG35il2bbjQqNNBIAUFwcD0mDJqeWxaAQkbaL_SGLHpGU6Ul6DGNpyOE4jHxASI7B8DdiwRJkuI-zsA-4uISC8kIkyskOjv-HGFi7mwsxwScbtsWQ5d6dK9xzNCjbwYCVrskJHjEV6aXIYI1VisgbNFyx-hsvbVKKydka1Olu6iurv25ZopDgU0LwRtppALYMURcZYRm9Nrh6ZHLq5306rqUXsKwBbmzKQ8WMExfICvvk1zRPbCNG6-wJCf0E0Ft1XNuFfhbDIM-2N_zS3QCJQGw"
|
|
|
|
async def test_chat_flow(message: str, location_tag: str = None):
|
|
print(f"\n=== Testing Gemma Tool Call Flow ===")
|
|
print(f"User Message: {message}")
|
|
print(f"Initial Location Tag: {location_tag}")
|
|
|
|
# 1. Setup Mock User
|
|
current_user = UserBase(
|
|
user_id=str(uuid.uuid4()),
|
|
name="Test User",
|
|
role="admin"
|
|
)
|
|
|
|
# 2. Setup DB Sessions
|
|
async with async_session() as db_session:
|
|
async with collector_async_session() as collector_db_session:
|
|
|
|
# 3. Prepare Conversation ID
|
|
conv_uuid = uuid.uuid4()
|
|
|
|
# Ensure Conversation exists in DB to prevent IntegrityError
|
|
new_conv = ConversationModel(
|
|
id=conv_uuid,
|
|
conversation_id=conv_uuid,
|
|
user_id=uuid.UUID(current_user.user_id),
|
|
conversation_title=f"Test Gemma Tool",
|
|
created_by=current_user.name,
|
|
updated_by=current_user.name
|
|
)
|
|
db_session.add(new_conv)
|
|
await db_session.commit()
|
|
print(f"✅ Created dummy conversation in DB: {conv_uuid}")
|
|
|
|
# 4. Create Request
|
|
request_as_rag = MagicMock()
|
|
request_as_rag.message = message
|
|
request_as_rag.conversation_id = conv_uuid
|
|
request_as_rag.attachments = []
|
|
request_as_rag.equipment_id = None
|
|
request_as_rag.model = "gemma4-e2b-llamacpp" # Force use of Gemma model
|
|
request_as_rag.context_state = {}
|
|
|
|
# Mocking CommonParameters
|
|
common = {
|
|
"db_session": db_session,
|
|
"page": 1,
|
|
"items_per_page": 5
|
|
}
|
|
|
|
# 5. Invoke stream_create
|
|
print("Starting stream...")
|
|
try:
|
|
generator = stream_create(
|
|
db_session=db_session,
|
|
collector_db_session=collector_db_session,
|
|
request=request_as_rag,
|
|
upload_type=None,
|
|
location_tag=location_tag,
|
|
current_user=current_user,
|
|
token=TOKEN,
|
|
common=common
|
|
)
|
|
|
|
full_response = ""
|
|
async for chunk in generator:
|
|
if chunk.startswith("data: "):
|
|
content = chunk[6:].strip()
|
|
if content:
|
|
if content.startswith("{"):
|
|
try:
|
|
data = json.loads(content)
|
|
print(f"[Metadata]: {content}")
|
|
except:
|
|
print(f"[Chunk]: {content}")
|
|
else:
|
|
# Normal text
|
|
print(content, end="", flush=True)
|
|
full_response += content
|
|
elif chunk.startswith("event: "):
|
|
print(f"\n[Event]: {chunk.strip()}")
|
|
|
|
print("\n\n=== Final Response ===")
|
|
print(full_response)
|
|
|
|
except Exception as e:
|
|
print(f"\n❌ Chat flow failed: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
async def main():
|
|
# Test the specific query that was failing
|
|
query = "Ambil data sensor PFI untuk HP FEEDWATER HEATER NO.5. Anda HARUS memanggil tool yang sesuai untuk mendapatkan data ini."
|
|
tag = "3FW-H050"
|
|
await test_chat_flow(query, location_tag=tag)
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|