cogview.py

import streamlit as st
from zhipuai import ZhipuAI
from zhipuai.types.image import GeneratedImage

from .config import COGVIEW_MODEL, ZHIPU_AI_KEY
from .interface import ToolObservation

@st.cache_resource
def get_zhipu_client():
    return ZhipuAI(api_key=ZHIPU_AI_KEY)

def map_response(img: GeneratedImage):
    return ToolObservation(
        content_type='image',
        text='CogView 已经生成并向用户展示了生成的图片。',
        image_url=img.url,
        role_metadata='cogview_result'
    )

def tool_call(prompt: str, session_id: str) -> list[ToolObservation]:
    client = get_zhipu_client()
    response = client.images.generations(model=COGVIEW_MODEL, prompt=prompt).data
    return list(map(map_response, response))