import uvicorn
import os
 
from typing import Annotated
from fastapi import (
    Depends,
    FastAPI,
    WebSocket,
    WebSocketException,
    WebSocketDisconnect,
    status,
)
from langchain.agents import create_structured_chat_agent, AgentExecutor
from langchain_community.utilities import SerpAPIWrapper
 
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
 
os.environ["OPENAI_API_KEY"] = 'sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx'    #你的Open AI Key
os.environ["SERPAPI_API_KEY"] = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
 
 
class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []
 
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
 
    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)
 
    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)
 
    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)
 
manager = ConnectionManager()
 
app = FastAPI()
 
async def authenticate(
    websocket: WebSocket,
    userid: str,
    secret: str,
):
    if userid is None or secret is None:
        raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
 
    print(f'userid: {userid},secret: {secret}')
    if '12345' == userid and 'xxxxxxxxxxxxxxxxxxxxxxxxxx' == secret:
        return 'pass'
    else:
        return 'fail'
 
@tool
def search(query:str):
    """只有需要了解实时信息或不知道的事情的时候才会使用这个工具，需要传入要搜索的内容。"""
    serp = SerpAPIWrapper()
    result = serp.run(query)
    print("实时搜索结果:", result)
    return result
 
 
def get_prompt():
    template='''
    Respond to the human as helpfully and accurately as possible. You have access to the following tools:
    
    {tools}
    
    Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).
    
    Valid "action" values: "Final Answer" or {tool_names}
    
    Provide only ONE action per $JSON_BLOB, as shown:
    
    ```
    
    {{
    
      "action": $TOOL_NAME,
    
      "action_input": $INPUT
    
    }}
    
    ```
    
    Follow this format:
    
    Question: input question to answer
    
    Thought: consider previous and subsequent steps
    
    Action:
    
    ```
    
    $JSON_BLOB
    
    ```
    
    Observation: action result
    
    ... (repeat Thought/Action/Observation N times)
    
    Thought: I know what to respond
    
    Action:
    
    ```
    
    {{
    
      "action": "Final Answer",
    
      "action_input": "Final response to human"
    
    }}
    
    Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation
    '''
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_template='''
    {input}
    
    {agent_scratchpad}
    
     (reminder to respond in a JSON blob no matter what)
    '''
    human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
    prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
 
    return prompt
 
async def chat(query):
    global llm,tools
    agent = create_structured_chat_agent(
        llm, tools, get_prompt()
    )
 
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True)
 
    result = agent_executor.invoke({"input": query})
    print(result['output'])
    yield result['output']
 
@app.websocket("/ws")
async def websocket_endpoint(*,websocket: WebSocket,userid: str,permission: Annotated[str, Depends(authenticate)],):
    await manager.connect(websocket)
    try:
        while True:
            text = await websocket.receive_text()
 
            if 'fail' == permission:
                await manager.send_personal_message(
                    f"authentication failed", websocket
                )
            else:
                if text is not None and len(text) > 0:
                    async for msg in chat(text):
                        await manager.send_personal_message(msg, websocket)
 
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        print(f"Client #{userid} left the chat")
        await manager.broadcast(f"Client #{userid} left the chat")
 
if __name__ == '__main__':
    tools = [search]
    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, max_tokens=512)
    uvicorn.run(app, host='0.0.0.0',port=7777)

    
        
    
    
        
WebSocket Chat
        
"" onsubmit="sendMessage(event)">
            USERID: <input type="text" id="userid" autocomplete="off" value="12345"/>
            SECRET: <input type="text" id="secret" autocomplete="off" value="xxxxxxxxxxxxxxxxxxxxxxxxxx"/>
            
            "connect(event)">Connect
            
            Message: <input type="text" id="messageText" autocomplete="off"/>
            Send
        
        
id='messages'>