215 lines
8.2 KiB
Python
215 lines
8.2 KiB
Python
import json
|
|
import re
|
|
|
|
import httpx
|
|
|
|
from channels.db import database_sync_to_async
|
|
from django.conf import settings
|
|
from django.utils import timezone
|
|
|
|
from apps.onboarding.consumers.base import BaseOnboardingConsumer, LogType
|
|
from apps.onboarding.consumers.prompts import OnboardingPrompts
|
|
from apps.onboarding.models import AgentInteractionLog, OnboardingSession
|
|
|
|
__all__ = ['OnboardingKnowledgeConsumer']
|
|
|
|
class OnboardingKnowledgeConsumer(BaseOnboardingConsumer):
|
|
"""
|
|
Route: /ws/onboarding/knowledge/<uuid:session_uuid>/
|
|
"""
|
|
|
|
session_uuid: str
|
|
|
|
def parse_extra(self):
|
|
self.session_uuid = self.scope['url_route']['kwargs'].get('session_uuid')
|
|
|
|
async def action_ask(self, data: dict):
|
|
page_uuid = data.get('page_uuid')
|
|
user_message = data.get('message')
|
|
mode = str(data.get('mode', 'separate'))
|
|
|
|
if not page_uuid or not user_message:
|
|
return await self.send_error('page_uuid and message are required.')
|
|
|
|
session = await self.get_session(self.session_uuid, self.user.id)
|
|
if not session:
|
|
return await self.send_error('Session not found or access denied.')
|
|
|
|
if not session.flow:
|
|
return await self.send_error('Onboarding flow not found.')
|
|
|
|
page = self._get_page(session.flow, str(page_uuid))
|
|
if not isinstance(page, dict):
|
|
return await self.send_error('Page not found in this flow.')
|
|
|
|
page_title = str(page.get('title') or 'Onboarding Page')
|
|
page_body = str(page.get('body') or '')
|
|
role_name = session.role.name
|
|
role_uuid = str(session.role.uuid)
|
|
|
|
config = await self.get_config_by_type(role_uuid, 'knowledge')
|
|
|
|
updated_page = False
|
|
revised_body = None
|
|
assistant_message = ''
|
|
|
|
if mode == 'update_page':
|
|
await self.send_log(LogType.STATUS, 'Revising page content...')
|
|
revised_body = await self._call_llm(
|
|
config,
|
|
OnboardingPrompts.ka_page_revision_prompt(role_name, page_title, page_body, str(user_message)),
|
|
max_tokens=3000,
|
|
stop=['\n[END]', '[END]'],
|
|
)
|
|
if revised_body:
|
|
await self.save_page_override(session, str(page_uuid), revised_body)
|
|
updated_page = True
|
|
assistant_message = (
|
|
'Updated this page by integrating your clarification request into the core content. '
|
|
'Please review the revised page text above.'
|
|
)
|
|
|
|
if not assistant_message:
|
|
await self.send_log(LogType.STATUS, 'Thinking...')
|
|
if config:
|
|
assistant_message = await self._call_llm(
|
|
config,
|
|
OnboardingPrompts.ka_help_prompt(role_name, page_title, page_body, str(user_message)),
|
|
max_tokens=1024,
|
|
) or OnboardingPrompts.KA_HELP_FALLBACK
|
|
else:
|
|
assistant_message = OnboardingPrompts.KA_HELP_FALLBACK
|
|
|
|
await self.save_page_help(session, str(page_uuid), str(user_message), assistant_message)
|
|
await self.log_interaction(session, str(user_message), assistant_message, str(page_uuid), mode, updated_page, config=config)
|
|
|
|
await self.send_log(LogType.COMPLETED, assistant_message, {
|
|
'updated_page': updated_page,
|
|
'revised_page_body': revised_body if mode == 'update_page' else None,
|
|
})
|
|
|
|
async def _call_llm(
|
|
self,
|
|
config,
|
|
prompt: str,
|
|
*,
|
|
max_tokens: int = 1024,
|
|
stop: list[str] | None = None,
|
|
) -> str | None:
|
|
if not config:
|
|
return None
|
|
system_prompt = config.system_prompt or OnboardingPrompts.FALLBACK_SYSTEM_PROMPT
|
|
llm_config = config.llm_config if isinstance(config.llm_config, dict) else {}
|
|
payload: dict = {
|
|
'model': llm_config.get('model_id', 'meta-llama-3.1-8b'),
|
|
'messages': [
|
|
{'role': 'system', 'content': system_prompt},
|
|
{'role': 'user', 'content': prompt},
|
|
],
|
|
'max_tokens': max_tokens,
|
|
'stream': True,
|
|
}
|
|
if stop:
|
|
payload['stop'] = stop
|
|
try:
|
|
chunks: list[str] = []
|
|
async with httpx.AsyncClient(timeout=120.0, auth=settings.INFERENCE_AUTH) as client:
|
|
async with client.stream('POST', settings.INFERENCE_CHAT_COMPLETIONS_ENDPOINT, json=payload) as response:
|
|
response.raise_for_status()
|
|
async for line in response.aiter_lines():
|
|
if not line.startswith('data: '):
|
|
continue
|
|
data = line[6:].strip()
|
|
if data == '[DONE]':
|
|
break
|
|
try:
|
|
chunk_obj = json.loads(data)
|
|
choice = chunk_obj['choices'][0]
|
|
delta = choice.get('delta', {}).get('content', '')
|
|
if delta:
|
|
chunks.append(delta)
|
|
await self.send_log(LogType.STREAM_CHUNK, delta)
|
|
if choice.get('finish_reason') == 'length':
|
|
self.logger.warning('Knowledge LLM response truncated (finish_reason=length)')
|
|
except Exception:
|
|
continue
|
|
result = ''.join(chunks).strip()
|
|
result = re.sub(r'\n?\[END\]\s*$', '', result).strip()
|
|
return result or None
|
|
except Exception as e:
|
|
self.logger.exception('Knowledge LLM call failed: %s', e)
|
|
return None
|
|
|
|
def _get_page(self, flow, page_uuid: str) -> dict | None:
|
|
pages = flow.structure if isinstance(flow.structure, list) else []
|
|
return next(
|
|
(p for p in pages if isinstance(p, dict) and str(p.get('uuid')) == page_uuid),
|
|
None,
|
|
)
|
|
|
|
@database_sync_to_async
|
|
def get_session(self, session_uuid: str, user_id: int):
|
|
return (
|
|
OnboardingSession.objects
|
|
.select_related('flow', 'role')
|
|
.filter(uuid=session_uuid, user_id=user_id)
|
|
.first()
|
|
)
|
|
|
|
@database_sync_to_async
|
|
def save_page_help(self, session, page_uuid: str, user_message: str, assistant_message: str):
|
|
state = session.state or {}
|
|
page_help = state.get('page_help', {})
|
|
if not isinstance(page_help, dict):
|
|
page_help = {}
|
|
|
|
thread = page_help.get(page_uuid, [])
|
|
if not isinstance(thread, list):
|
|
thread = []
|
|
|
|
thread.append({
|
|
'question': user_message,
|
|
'answer': assistant_message,
|
|
'timestamp': timezone.now().isoformat(),
|
|
})
|
|
|
|
page_help[page_uuid] = thread[-20:]
|
|
state['page_help'] = page_help
|
|
session.state = state
|
|
session.save(update_fields=['state', 'updated_at'])
|
|
|
|
@database_sync_to_async
|
|
def save_page_override(self, session, page_uuid: str, new_body: str):
|
|
state = session.state if isinstance(session.state, dict) else {}
|
|
overrides = state.get('page_overrides', {})
|
|
if not isinstance(overrides, dict):
|
|
overrides = {}
|
|
overrides[page_uuid] = new_body
|
|
state['page_overrides'] = overrides
|
|
session.state = state
|
|
session.save(update_fields=['state', 'updated_at'])
|
|
|
|
@database_sync_to_async
|
|
def log_interaction(
|
|
self, session, user_message: str, assistant_message: str,
|
|
page_uuid: str, mode: str, updated_page: bool, config=None,
|
|
):
|
|
AgentInteractionLog.objects.create(
|
|
session=session,
|
|
agent_config=config,
|
|
sender_type='user',
|
|
content=user_message,
|
|
tool_call_metadata={'action': 'ask_ka', 'page_uuid': page_uuid, 'mode': mode},
|
|
)
|
|
AgentInteractionLog.objects.create(
|
|
session=session,
|
|
agent_config=config,
|
|
sender_type='ai',
|
|
content=assistant_message,
|
|
tool_call_metadata={
|
|
'action': 'ask_ka_response',
|
|
'page_uuid': page_uuid,
|
|
'mode': mode,
|
|
'updated_page': updated_page,
|
|
},
|
|
)
|