AIDC-AI-ComfyUI-Copilot/service/conversation_service.py

# Copyright (C) 2025 AIDC-AI
# Licensed under the MIT License.

import json
import os
import asyncio
import time
from typing import Optional, Dict, Any, TypedDict, List, Union

import server
from aiohttp import web
import aiohttp
import base64

# 使用内存字典存储会话消息
session_messages = {}

# 在文件开头添加
STATIC_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "public")

# Define types using TypedDict
class Node(TypedDict):
    name: str  # 节点名称
    description: str  # 节点描述
    image: str  # 节点图片url，可为空
    github_url: str  # 节点github地址
    from_index: int  # 节点在列表中的位置
    to_index: int  # 节点在列表中的位置

class NodeInfo(TypedDict):
    existing_nodes: List[Node]  # 已安装的节点
    missing_nodes: List[Node]  # 未安装的节点

class Workflow(TypedDict, total=False):
    id: Optional[int]  # 工作流id
    name: Optional[str]  # 工作流名称
    description: Optional[str]  # 工作流描述
    image: Optional[str]  # 工作流图片
    workflow: Optional[str]  # 工作流

class ExtItem(TypedDict):
    type: str  # 扩展类型
    data: Union[dict, list]  # 扩展数据

class ChatResponse(TypedDict):
    session_id: str  # 会话id
    text: Optional[str]  # 返回文本
    finished: bool  # 是否结束
    type: str  # 返回的类型
    format: str  # 返回的格式
    ext: Optional[List[ExtItem]]  # 扩展信息

def get_workflow_templates():
    templates = []
    workflows_dir = os.path.join(STATIC_DIR, "workflows")
    
    for filename in os.listdir(workflows_dir):
        if filename.endswith('.json'):
            with open(os.path.join(workflows_dir, filename), 'r') as f:
                template = json.load(f)
                templates.append(template)
    
    return templates

@server.PromptServer.instance.routes.get("/workspace/fetch_messages_by_id")
async def fetch_messages(request):
    session_id = request.query.get('session_id')
    data = await asyncio.to_thread(fetch_messages_sync, session_id)
    return web.json_response(data)

def fetch_messages_sync(session_id):
    print("fetch_messages: ", session_id)
    return session_messages.get(session_id, [])

@server.PromptServer.instance.routes.post("/workspace/workflow_gen")
async def workflow_gen(request):
    print("Received workflow_gen request")
    req_json = await request.json()
    print("Request JSON:", req_json)
    
    response = web.StreamResponse(
        status=200,
        reason='OK',
        headers={
            'Content-Type': 'application/json',
            'X-Content-Type-Options': 'nosniff'
        }
    )
    await response.prepare(request)
    
    session_id = req_json.get('session_id')
    user_message = req_json.get('message')
    
    # Create user message
    user_msg = {
        "id": str(len(session_messages.get(session_id, []))),
        "content": user_message,
        "role": "user"
    }
    
    if "workflow" in user_message.lower():
        workflow = {
            "name": "basic_image_gen",
            "description": "Create a basic image generation workflow",
            "image": "https://placehold.co/600x400",
            "workflow": """{ ... }"""  # Your workflow JSON here
        }
        
        chat_response = ChatResponse(
            session_id=session_id,
            text="",
            finished=False,
            type="workflow_option",
            format="text",
            ext=[{"type": "workflows", "data": [workflow]}]
        )
        
        await response.write(json.dumps(chat_response).encode() + b"\n")
        
        message = "Let me help you choose a workflow. Here are some options available:"
        accumulated = ""
        for char in message:
            accumulated += char
            chat_response["text"] = accumulated
            await response.write(json.dumps(chat_response).encode() + b"\n")
            await asyncio.sleep(0.01)
        
        chat_response["finished"] = True
        chat_response["text"] = message
        await response.write(json.dumps(chat_response).encode() + b"\n")
        
    elif "recommend" in user_message.lower():
        existing_nodes = [
            {
                "name": "LoraLoader",
                "description": "Load LoRA weights for conditioning.",
                "image": "",
                "github_url": "https://github.com/CompVis/taming-transformers",
                "from_index": 0,
                "to_index": 0
            },
            {
                "name": "KSampler",
                "description": "Generate images using K-diffusion sampling.",
                "image": "",
                "github_url": "https://github.com/CompVis/taming-transformers",
                "from_index": 0,
                "to_index": 0
            }
        ]
        
        missing_nodes = [
            {
                "name": "CLIPTextEncode",
                "description": "Encode text prompts for conditioning.",
                "image": "",
                "github_url": "https://github.com/CompVis/clip-interrogator",
                "from_index": 0,
                "to_index": 0
            }
        ]
        
        node_info = {
            "existing_nodes": existing_nodes,
            "missing_nodes": missing_nodes
        }
        
        chat_response = ChatResponse(
            session_id=session_id,
            text="",
            finished=False,
            type="downstream_node_recommend",
            format="text",
            ext=[{"type": "node_info", "data": node_info}]
        )
        
        await response.write(json.dumps(chat_response).encode() + b"\n")
        
        message = "Here are some recommended nodes:"
        accumulated = ""
        for char in message:
            accumulated += char
            chat_response["text"] = accumulated
            await response.write(json.dumps(chat_response).encode() + b"\n")
            await asyncio.sleep(0.01)
        
        chat_response["finished"] = True
        chat_response["text"] = message
        await response.write(json.dumps(chat_response).encode() + b"\n")
        
    else:
        chat_response = ChatResponse(
            session_id=session_id,
            text="",
            finished=False,
            type="message",
            format="text",
            ext=[{"type": "guides", "data": ["Create a workflow", "Search for nodes", "Get node recommendations"]}]
        )
        
        await response.write(json.dumps(chat_response).encode() + b"\n")
        
        message = "I can help you with workflows, nodes, and more. Try asking about:"
        accumulated = ""
        for char in message:
            accumulated += char
            chat_response["text"] = accumulated
            await response.write(json.dumps(chat_response).encode() + b"\n")
            await asyncio.sleep(0.01)
        
        chat_response["finished"] = True
        chat_response["text"] = message
        await response.write(json.dumps(chat_response).encode() + b"\n")
    
    if session_id not in session_messages:
        session_messages[session_id] = []
    
    session_messages[session_id].extend([user_msg])
    
    await response.write_eof()
    return response

async def upload_to_oss(file_data: bytes, filename: str) -> str:
    # Implement your OSS upload logic here
    # Return the URL of the uploaded file
    pass

@server.PromptServer.instance.routes.post("/api/chat/invoke")
async def invoke_chat(request):
    data = await request.json()
    session_id = data['session_id']
    prompt = data['prompt']
    images = data.get('images', [])

    image_urls = []
    for image in images:
        # 从base64解码图片数据
        image_data = base64.b64decode(image['data'].split(',')[1])
        # 上传到OSS
        url = await upload_to_oss(image_data, image['filename'])
        image_urls.append(url)

    # 使用image_urls进行后续处理...