46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
import tyro
|
|
from typing import Annotated
|
|
import uuid
|
|
from loguru import logger
|
|
|
|
from lang_agent.pipeline import Pipeline, PipelineConfig
|
|
from lang_agent.config import load_tyro_conf
|
|
from lang_agent.components.conv_store import use_printer
|
|
|
|
|
|
def main(
|
|
conf: PipelineConfig,
|
|
stream: Annotated[bool, tyro.conf.arg(name="stream")] = True,
|
|
):
|
|
"""Demo chat script for langchain-agent pipeline.
|
|
|
|
Args:
|
|
conf: Pipeline configuration
|
|
stream: Enable streaming mode for chat responses
|
|
"""
|
|
use_printer()
|
|
if conf.config_f is not None:
|
|
conf = load_tyro_conf(conf.config_f)
|
|
|
|
logger.info(conf)
|
|
pipeline: Pipeline = conf.setup()
|
|
thread_id = str(uuid.uuid4())
|
|
while True:
|
|
user_input = input("请讲：")
|
|
if user_input.lower() == "exit":
|
|
break
|
|
|
|
if stream:
|
|
# Streaming mode: print chunks as they arrive
|
|
print("回答: ", end="", flush=True)
|
|
for chunk in pipeline.chat(user_input, as_stream=True, thread_id=thread_id):
|
|
print(chunk, end="", flush=True)
|
|
print() # New line after streaming completes
|
|
else:
|
|
# Non-streaming mode: print full response
|
|
response = pipeline.chat(user_input, as_stream=False, thread_id=thread_id)
|
|
print(f"回答: {response}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tyro.cli(main) |