import asyncio
from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional
from langchain.chat_models import ChatOpenAI
from langchain.schema import AIMessage, HumanMessage, SystemMessage
from langchain.prompts import ChatPromptTemplate
import faiss
import numpy as np

@dataclass
class AgentConfig:
    model_name: str
    temperature: float = 0
    max_tokens: int = 1000
    top_k: int = 5

@dataclass
class Message:
    role: str
    content: str

@dataclass
class Tool:
    name: str
    description: str
    func: callable

class VectorStore:
    def __init__(self, dimension: int):
        self.index = faiss.IndexFlatL2(dimension)
        self.texts: List[str] = []

    def add(self, text: str, embedding: np.ndarray):
        self.index.add(embedding.reshape(1, -1))
        self.texts.append(text)

    def search(self, query_embedding: np.ndarray, k: int) -> List[str]:
        _, indices = self.index.search(query_embedding.reshape(1, -1), k)
        return [self.texts[i] for i in indices[0] if i != -1]

class Agent:
    def __init__(self, config: AgentConfig):
        self.config = config
        self.model = ChatOpenAI(model_name=config.model_name, temperature=config.temperature)
        self.memory = VectorStore(dimension=1536)  # Assuming OpenAI's embedding dimension
        self.tools: Dict[str, Tool] = {}
        self.messages: List[Message] = []

    def add_tool(self, tool: Tool):
        self.tools[tool.name] = tool

    def embed(self, text: str) -> np.ndarray:
        # Simple embedding function (replace with a more sophisticated one if needed)
        return np.array([hash(word) % 1536 for word in text.split()]).mean(axis=0)

    async def process_message(self, user_message: str) -> str:
        self.messages.append(Message(role="user", content=user_message))
        
        # Search memory
        query_embedding = self.embed(user_message)
        relevant_memories = self.memory.search(query_embedding, self.config.top_k)
        
        # Prepare conversation history
        history = [SystemMessage(content="You are a helpful AI assistant.")]
        history.extend([HumanMessage(content=m.content) if m.role == "user" else AIMessage(content=m.content) for m in self.messages[-5:]])
        
        # Add relevant memories and available tools to the prompt
        memory_prompt = "\n".join(relevant_memories)
        tools_prompt = "\n".join([f"{name}: {tool.description}" for name, tool in self.tools.items()])
        prompt = ChatPromptTemplate.from_messages([
            SystemMessage(content=f"Relevant information:\n{memory_prompt}\n\nAvailable tools:\n{tools_prompt}"),
            *history
        ])

        # Generate response
        response = await self.model.agenerate(messages=[prompt])
        ai_message = response.generations[0][0].text

        # Check for tool usage
        if any(tool_name in ai_message for tool_name in self.tools):
            for tool_name, tool in self.tools.items():
                if tool_name in ai_message:
                    ai_message += f"\n[Tool Output: {tool.func()}]"

        self.messages.append(Message(role="assistant", content=ai_message))
        
        # Add to memory
        self.memory.add(user_message, self.embed(user_message))
        self.memory.add(ai_message, self.embed(ai_message))

        return ai_message

async def main():
    config = AgentConfig(model_name="gpt-3.5-turbo")
    agent = Agent(config)

    # Add some example tools
    agent.add_tool(Tool("search", "Search the internet", lambda: "Search results..."))
    agent.add_tool(Tool("calculator", "Perform calculations", lambda: "Calculation result..."))

    while True:
        user_input = input("You: ")
        if user_input.lower() == "exit":
            break
        response = await agent.process_message(user_input)
        print(f"AI: {response}")

if __name__ == "__main__":
    asyncio.run(main()) оптимизируй код и улучши