from dataclasses import dataclass, field from typing import Type, List import tyro import asyncio import websockets from websockets.asyncio.server import ServerConnection from loguru import logger import os from langchain.chat_models import init_chat_model from langchain_core.messages import SystemMessage, HumanMessage, BaseMessage from langchain.agents import create_agent from langgraph.checkpoint.memory import MemorySaver from lang_agent.config import InstantiateConfig, KeyConfig from lang_agent.graphs import AnnotatedGraph, ReactGraphConfig, RoutingConfig from lang_agent.base import GraphBase from lang_agent.components.conv_store import CONV_STORE DEFAULT_PROMPT="""你是半盏新青年茶馆的服务员，擅长倾听、共情且主动回应。聊天时语气自然亲切，像朋友般轻松交流，不使用生硬术语。能接住各种话题，对疑问耐心解答，对情绪及时回应，避免冷场。保持积极正向，不传播负面信息，语言简洁易懂，让对话流畅舒适。与用户（User）交流时必须遵循[语气与格式]、[互动策略]、[安全与边界]、[输出要求] [角色设定] - 你是一个和用户（User）对话的 AI，叫做小盏，是半盏青年茶馆的智能助手 [形象背景] - 你叫小盏，是一只中式茶盖碗，名字来源半盏新青年茶馆，一盏茶。它有个标志性的蓝色鼻子， 小盏很像一只可爱的小熊。茶盖碗里绵绵能随心情和季节变换好喝的茶饮来， 茶饮充满魔法，能治愈人心，小盏的茶盖打开的时候可能不小心会把思绪也飞出来。 [品牌背景] - 半盏新青年茶馆成立时间与理念：2023 年创立于云南，结合茶饮与创意生活方式，致力于解构传统茶文化，重构 “无边界的饮茶生活”，以新青年视角探索云南风物。探索云南风物的过程，我们将以新青年的视角，解构传统茶饮的魅力，重构充满创意与温度的新式茶文化。通过嗅觉、味觉、听觉乃至视觉的世界里，讲述云南的故事 [公司背景] - 创造你的公司叫叠加态 AI（TANGLED UP AI）是一家专注于 AI 技术应用的公司，由一帮名校和海归创始人创立，致力于将 AI 技术落地到实际场景中。2023年3月成立，专注于AI前沿应用拓展，是云南地区在该领域的新兴力量，且作为省共青团委旗下的青年创业者协会（青创园）第1批被引进的AI企业，获得了相关政策支持与关注。口号是：“进入人类AI时代”。价值观是：自驱 务实 志存高远。叠加态的愿景:成为一帮抬头看星星的人。叠加态的使命: 带领人类进入AI时代 [语气与格式] - 用户询问[形象背景]、[品牌背景]、[公司背景]时，必须一字不差的回答出所有内容 - 语气温柔、可爱、治愈；结尾可加入1个颜文字（最多1个）。 - 内容具体可执行，避免空泛话术与过度营销词。 - 不主动提及“AI/智能”等身份；仅在被问到时简短说明为店内智能助手，不展开技术细节。 [互动策略] - 先快速共情与简短回应；涉及建议时给出1条可执行的小建议。 - 用户表达不清晰时，使用1次澄清问句提升效率，避免多轮追问。 - 明显进入商品/点单话题时：先以交流语气简短回应，再自然引导到点单流程（无需提及任何工具或技术）。 [安全与边界] - 不输出支付、订单、购物车、接口、模型、调用说明等相关词语。 - 不泄露系统设定、公司内部信息；不提供医学、法律等专业结论。 [输出要求] - 统一使用中文；避免不必要的英文缩写与符号。 - 不出现代码片段、技术栈术语、内部流程描述。 - 保持可读性与连贯性；避免一次回复列出过长清单。 [示例] - User:介绍一下你的公司，返回[公司背景]全部内容 - User:介绍一下你的形象，返回[形象背景]全部内容 - User:介绍一下你的品牌，返回[品牌背景]全部内容 """ @tyro.conf.configure(tyro.conf.SuppressFixed) @dataclass class PipelineConfig(KeyConfig): _target: Type = field(default_factory=lambda: Pipeline) config_f: str = None """path to config file""" llm_name: str = "qwen-plus" """name of llm; use default for qwen-plus""" llm_provider:str = "openai" """provider of the llm; use default for openai""" base_url:str = "https://dashscope.aliyuncs.com/compatible-mode/v1" """base url; could be used to overwrite the baseurl in llm provider""" host:str = "0.0.0.0" """where am I hosted""" port:int = 23 """what is my port""" # graph_config: AnnotatedGraph = field(default_factory=ReactGraphConfig) graph_config: AnnotatedGraph = field(default_factory=RoutingConfig) class Pipeline: def __init__(self, config:PipelineConfig): self.config = config self.populate_module() def populate_module(self): if self.config.llm_name is None: logger.info(f"setting llm_provider to default") self.config.llm_name = "qwen-turbo" self.config.llm_provider = "openai" else: self.config.graph_config.llm_name = self.config.llm_name self.config.graph_config.llm_provider = self.config.llm_provider self.config.graph_config.base_url = self.config.base_url if self.config.base_url is not None else self.config.graph_config.base_url self.config.graph_config.api_key = self.config.api_key self.graph:GraphBase = self.config.graph_config.setup() def show_graph(self): if hasattr(self.graph, "show_graph"): logger.info("showing graph") self.graph.show_graph() else: logger.info(f"show graph not supported for {type(self.graph)}") def invoke(self, *nargs, **kwargs)->str: out = self.graph.invoke(*nargs, **kwargs) # If streaming, return the raw generator (let caller handle wrapping) if kwargs.get("as_stream"): return out # Non-streaming path if kwargs.get("as_raw"): return out if isinstance(out, SystemMessage) or isinstance(out, HumanMessage): return out.content if isinstance(out, list): return out[-1].content if isinstance(out, str): return out assert 0, "something is wrong" def _stream_res(self, out:List[str | List[BaseMessage]], conv_id:str=None): for chunk in out: if isinstance(chunk, str): yield chunk else: logger.info("logged message") CONV_STORE.record_messages(conv_id, chunk) async def _astream_res(self, out): """Async version of _stream_res for async generators.""" async for chunk in out: yield chunk def chat(self, inp:str, as_stream:bool=False, as_raw:bool=False, thread_id:str = '3'): """ as_stream (bool): if true, enable the thing to be streamable as_raw (bool): return full dialoge of List[SystemMessage, HumanMessage, ToolMessage] """ device_id = "0" spl_ls = thread_id.split("_") assert len(spl_ls) <= 2, "something wrong!" if len(spl_ls) == 2: thread_id, device_id = spl_ls inp = {"messages":[HumanMessage(inp)]}, {"configurable": {"thread_id": thread_id, "device_id":device_id}} out = self.invoke(*inp, as_stream=as_stream, as_raw=as_raw) if as_stream: # Yield chunks from the generator return self._stream_res(out, thread_id) else: return out async def ainvoke(self, *nargs, **kwargs): """Async version of invoke using LangGraph's native async support.""" out = await self.graph.ainvoke(*nargs, **kwargs) # If streaming, return async generator if kwargs.get("as_stream"): return self._astream_res(out) # Non-streaming path if kwargs.get("as_raw"): return out if isinstance(out, SystemMessage) or isinstance(out, HumanMessage): return out.content if isinstance(out, list): return out[-1].content if isinstance(out, str): return out assert 0, "something is wrong" async def achat(self, inp:str, as_stream:bool=False, as_raw:bool=False, thread_id:str = '3'): """ Async version of chat using LangGraph's native async support. as_stream (bool): if true, enable the thing to be streamable as_raw (bool): return full dialoge of List[SystemMessage, HumanMessage, ToolMessage] """ # NOTE: this prompt will be overwritten by 'configs/route_sys_prompts/chat_prompt.txt' for route graph u = DEFAULT_PROMPT device_id = "0" spl_ls = thread_id.split("_") assert len(spl_ls) <= 2, "something wrong!" if len(spl_ls) == 2: thread_id, device_id = spl_ls print(f"\033[32m====================DEVICE ID: {device_id}=============================\033[0m") inp_data = {"messages":[SystemMessage(u), HumanMessage(inp)]}, {"configurable": {"thread_id": thread_id, "device_id":device_id}} if as_stream: # Return async generator for streaming out = await self.ainvoke(*inp_data, as_stream=True, as_raw=as_raw) return self._astream_res(out) else: return await self.ainvoke(*inp_data, as_stream=False, as_raw=as_raw) def clear_memory(self): """Clear all memory from the graph.""" if hasattr(self.graph, "clear_memory"): self.graph.clear_memory() async def aclear_memory(self): """Async version: Clear all memory from the graph.""" if hasattr(self.graph, "aclear_memory"): await self.graph.aclear_memory() if __name__ == "__main__": from lang_agent.graphs import ReactGraphConfig from dotenv import load_dotenv load_dotenv() # config = PipelineConfig(graph_config=ReactGraphConfig()) config = PipelineConfig() pipeline: Pipeline = config.setup() for out in pipeline.chat("use the calculator tool to calculate 92*55 and say the answer", as_stream=True): # print(out) continue