sql_demo/111.py

import requests
from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel
import aiohttp
import logging
from fastapi.responses import JSONResponse
import json
from typing import Optional, List, Dict
import asyncio

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def call_llm(question: str, userId: str) -> Optional[Dict]:
    url = "http://192.168.100.100:7861/chat/kb_chat"
    payload = {
        "query": question,
        "mode": "local_kb",
        "kb_name": "lqbz",
        "top_k": 3,
        "score_threshold": 2,
        "history": [
            {"content": '你好', "role": "user"},
            {"content": "", "role": "assistant"}
        ],
        "stream": True,  # 设置为True以启用流式传输
        "model": "deepseek-r1:1.5b",
        "temperature": 0.7,
        "max_tokens": 0,
        "prompt_name": "default",
        "return_direct": False
    }
    combined_text = ""
    all_docs = []
    try:
        with requests.post(url, json=payload, stream=True) as response:
            if response.status_code == 200:
                for line in response.iter_lines():
                    if line:  # 过滤掉keep-alive新行
                        decoded_line = line.decode('utf-8')
                        if decoded_line.startswith('data: '):
                            json_str = decoded_line[len('data: '):]
                            try:
                                json_obj = json.loads(json_str)

                                # 提取choices中的内容并合并
                                for choice in json_obj.get("choices", []):
                                    content = choice.get("delta", {}).get("content", "")
                                    if content:
                                        combined_text += content  # 直接追加，保留换行符

                                # 收集docs中的内容
                                docs = json_obj.get("docs", [])
                                all_docs.extend(docs)
                            except json.JSONDecodeError as e:
                                logger.error(f"无法解析的块: {json_str[:50]}...")
                                logger.error(f"错误信息: {e}")
                # print(combined_text)
                print("************")
                think = combined_text.split("</think>")[0]
                think = think.replace("<think>", "")
                print("think", think)
                answer = combined_text.split("</think>")[1]
                print("answer:", answer)
                print("************")

                # 返回结果，包括保留换行符的answer
                return {"answer": combined_text, "docs": all_docs}
            else:
                logger.error(f"Error calling service: {response.status_code} - {response.text}")
                return None
    except requests.RequestException as e:
        logger.error(f"Unexpected error while calling service: {e}")
        return None


# 示例：如何调用这个函数，并显示保留换行的答案
if __name__ == "__main__":
    question = "你好"
    userId = "用户ID"
    result = call_llm(question, userId)
    if result is not None:
        print("成功获取结果:")
        print("Answer:\n{}".format(result))  # 直接打印，保留换行符
        if result["docs"]:
            print("Docs:\n{}".format(result["docs"]))
    else:
        print("未能成功获取结果")