Python实现本地AnythingLLM对话

import sys
import subprocess
import importlib.util

# 检查必要的库是否已安装
required_packages = ['requests']
for package in required_packages:
    if importlib.util.find_spec(package) is None:
        print(f"正在安装必要的库: {package}")
        subprocess.check_call([sys.executable, "-m", "pip", "install", package])


import requests
import json
import base64
from typing import List, Optional, Dict, Any
import os

class AnythingLLMClient:
    def __init__(self, base_url: str, api_key: str, workspace_id: str):
        """
        初始化AnythingLLM客户端
        
        Args:
            base_url: AnythingLLM服务的基础URL
            api_key: API密钥
            workspace_id: 工作区ID
        """
        self.base_url = base_url
        self.api_key = api_key
        self.workspace_id = workspace_id
        self.headers = {
            'accept': 'application/json',
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }
    
    def chat(self, message: str, session_id: str = "default-session", 
             mode: str = "chat", attachments: Optional[List[Dict[str, str]]] = None) -> Dict[str, Any]:
        """
        发送聊天消息到AnythingLLM
        
        Args:
            message: 用户消息
            session_id: 会话ID，用于区分不同的对话
            mode: 对话模式，默认为"chat"
            attachments: 附件列表，如图片等
            
        Returns:
            API响应的JSON数据
        """
        url = f"{self.base_url}/api/v1/workspace/{self.workspace_id}/chat"
        
        payload = {
            "message": message,
            "mode": mode,
            "sessionId": session_id,
            "attachments": attachments or []
        }
        
        try:
            response = requests.post(url, headers=self.headers, json=payload)
            response.raise_for_status()  # 如果响应状态码不是200，将引发异常
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"请求错误: {e}")
            return {"error": str(e)}
    
    def add_image_attachment(self, image_path: str) -> Dict[str, str]:
        """
        将图片添加为附件
        
        Args:
            image_path: 图片文件路径
            
        Returns:
            包含图片信息的字典
        """
        try:
            with open(image_path, "rb") as image_file:
                encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
            
            # 获取文件扩展名并确定MIME类型
            _, ext = os.path.splitext(image_path)
            mime_type = f"image/{ext[1:]}" if ext[1:] else "image/png"
            
            return {
                "name": os.path.basename(image_path),
                "mime": mime_type,
                "contentString": f"data:{mime_type};base64,{encoded_string}"
            }
        except Exception as e:
            print(f"处理图片时出错: {e}")
            return {}


def main():
    # 配置参数
    base_url = "http://localhost:3001"
    api_key = ""  # 替换为你的API密钥
    workspace_id = "532b7ee0-3f6a-4f8c-982d-2fcae7d4597b"
    
    # 创建客户端
    client = AnythingLLMClient(base_url, api_key, workspace_id)
    
    print("输入'exit'退出。")
    
    session_id = "python-client-session-" + str(hash(api_key))[:8]
    
    while True:
        user_input = input("\n你: ")
        
        if user_input.lower() == 'exit':
            print("再见！")
            break
        
        attachments = []
        message = user_input
        
        # 检查是否包含图片命令
        if user_input.startswith("image:"):
            parts = user_input.split(" ", 1)
            image_path = parts[0][6:]  # 去掉"image:"前缀
            
            if len(parts) > 1:
                message = parts[1]
            else:
                message = "请分析这张图片"
                
            image_attachment = client.add_image_attachment(image_path)
            if image_attachment:
                attachments.append(image_attachment)
                print(f"已添加图片: {image_path}")
            else:
                print(f"无法添加图片: {image_path}")
                continue
        
        # 发送消息
        response = client.chat(message, session_id=session_id, attachments=attachments)
        
        # 处理响应
        if "error" in response and response["error"]:
            print(f"\n错误: {response['error']}")
        else:
            print("\nAnythingLLM: ", end="")
            if "textResponse" in response:
                print(response["textResponse"])
            else:
                print("无响应或响应格式不正确")
            
            # 显示来源信息
            if "sources" in response and response["sources"]:
                print("\n来源:")
                for source in response["sources"]:
                    print(f"- {source.get('title', '未知标题')}")


if __name__ == "__main__":
    main()