LangChain AI-агент на Python

# Было с OpenAI API напрямую:
def get_completion(prompt, history=[]):
    formatted_prompt = format_complex_prompt(prompt, history)
    try:
        response = openai.Completion.create(
            model="text-davinci-003",
            prompt=formatted_prompt,
            max_tokens=500
        )
        return parse_response(response.choices[0].text)
    except Exception as e:
        return handle_errors(e)
        
# Стало с LangChain:
from langchain.chains import ConversationChain
from langchain.llms import OpenAI
 
conversation = ConversationChain(
    llm=OpenAI(model_name="text-davinci-003"),
    verbose=True
)
response = conversation.predict(input="Привет, мир!")

from langchain_core.prompts import ChatPromptTemplate
 
# Пример продвинутого промпта для агента
system_prompt = """Ты умный асистент со следующими инструментами: {tools}
Используй JSON для вызова инструмента, указывая "action" и "action_input".
Допустимые значения для "action": "Final Answer" или {tool_names}
 
Формат должен быть таким:
{{
"action": НАЗВАНИЕ_ИНСТРУМЕНТА,
"action_input": ВХОДНЫЕ_ДАННЫЕ
}}
"""
 
human_prompt = "{input}\n{agent_scratchpad}"
 
prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", human_prompt),
])

from langchain_core.tools import tool
 
@tool
def get_todays_date() -> str:
    """Получает текущую дату в формате ГГГГ-ММ-ДД."""
    from datetime import datetime
    date = datetime.now().strftime("%Y-%m-%d")
    return date

# Пример Few-shot промпта для разбора дат
few_shot_template = """
Вопрос: Какой день недели был 15 мая 2023?
Мысль: Мне нужно определить день недели для конкретной даты.
Действие: Определяю день недели для 15 мая 2023...
Наблюдение: 15 мая 2023 был понедельник.
Ответ: 15 мая 2023 был понедельник.
 
Вопрос: {question}
"""

# Пример использования Plan-and-Execute
from langchain.agents import create_plan_and_execute_agent
 
agent = create_plan_and_execute_agent(
    llm=model,
    tools=tools,
    verbose=True
)
 
agent.invoke({
    "input": "Сравни погоду в Москве и Санкт-Петербурге на завтра"
})

# Пример реализации MultiAgent-паттерна
researcher = create_react_agent(model, research_tools, "Исследователь")
analyst = create_react_agent(model, analysis_tools, "Аналитик")
writer = create_react_agent(model, writing_tools, "Писатель")
 
def solve_complex_task(query):
    research_results = researcher.run(query)
    analysis = analyst.run(research_results)
    final_report = writer.run(f"Составь отчет на основе: {analysis}")
    return final_report

# Упрощенная схема выполнения цепочки
def invoke(self, input_data):
    # 1. Подготовка входных данных
    prepared_input = self._prepare_input(input_data)
    
    # 2. Основное выполнение
    try:
        raw_output = self._execute(prepared_input)
    except Exception as e:
        return self._handle_error(e)
    
    # 3. Постобработка результата
    return self._process_output(raw_output)

# Старый императивный подход
def process_query(query):
    formatted_prompt = prompt_template.format(query=query)
    llm_result = llm.generate(formatted_prompt)
    parsed_result = output_parser.parse(llm_result)
    return parsed_result
 
# Новый функциональный стиль
chain = prompt_template | llm | output_parser
result = chain.invoke({"query": "Какая сегодня погода?"})

# Условное ветвление в функциональном стиле
def route_to_correct_tool(input_data):
    if "погода" in input_data["query"].lower():
        return {"route": "weather"}
    else:
        return {"route": "general"}
 
chain = (
    prompt_template
    | llm
    | RunnablePassthrough.assign(route=route_to_correct_tool)
    | RunnableBranch(
        (lambda x: x["route"] == "weather", weather_chain),
        (lambda x: x["route"] == "general", general_chain),
    )
)

from langchain.callbacks import RunLoggerHandler
from langchain.globals import set_handler
 
# Добавляем глобальный обработчик для логирования
logger = RunLoggerHandler()
set_handler(logger)
 
# Теперь все цепочки будут логироваться автоматически
response = agent_executor.invoke({"input": "Что такое ИИ?"})
 
# Посмотрим, что записалось в лог
print(logger.get_runs())

class TokenCounterMiddleware:
    def __init__(self):
        self.total_tokens = 0
        
    def __call__(self, runnable, input_data, config, **kwargs):
        # Запускаем оригинальный runnable
        output = runnable.invoke(input_data, config)
        
        # Подсчитываем токены, если это вызов LLM
        if hasattr(output, "token_usage"):
            self.total_tokens += output.token_usage.total_tokens
            print(f"Использовано токенов: {self.total_tokens}")
            
        return output
 
# Подключаем наш middleware
counter = TokenCounterMiddleware()
chain_with_counter = chain.with_handlers([counter])

class CheckpointMiddleware:
    def __call__(self, runnable, input_data, config, **kwargs):
        # Проверяем, есть ли сохраненный чекпоинт
        checkpoint_key = hash(str(input_data))
        if checkpoint_key in self.checkpoints:
            print("Восстанавливаем из чекпоинта")
            return self.checkpoints[checkpoint_key]
            
        # Выполняем оригинальный вызов
        output = runnable.invoke(input_data, config)
        
        # Сохраняем результат
        self.checkpoints[checkpoint_key] = output
        return output

from langchain_core.tools import tool
 
@tool
def search_products(query: str) -> list:
    """Ищет товары в базе данных по запросу."""
    # Здесь может быть запрос к вашей базе данных
    products = database.search(query)
    return products

@tool
def calculate_delivery_time(postal_code: str, weight: float) -> str:
    """Рассчитывает время доставки на основе почтового индекса и веса посылки в кг."""
    # Логика расчета времени доставки
    if weight > 10:
        return "3-5 рабочих дней"
    return "1-2 рабочих дня"

@tool(return_direct=True)
def generate_image(prompt: str) -> str:
    """Генерирует изображение по текстовому описанию и возвращает URL."""
    # Вызов API генерации изображений
    image_url = image_generator.create(prompt)
    return image_url

from langchain.tools import BaseTool
 
class DatabaseQueryTool(BaseTool):
    name = "database_query"
    description = "Выполняет SQL-запрос к базе данных и возвращает результаты."
    
    def _run(self, query: str) -> list:
        try:
            connection = get_db_connection()
            cursor = connection.cursor()
            cursor.execute(query)
            results = cursor.fetchall()
            return results
        except Exception as e:
            return f"Ошибка выполнения запроса: {str(e)}"
            
    async def _arun(self, query: str) -> list:
        # Асинхронная версия метода
        return await async_db_query(query)

tools = [search_products, calculate_delivery_time, generate_image, DatabaseQueryTool()]
 
agent_executor = AgentExecutor(
    agent=chain,
    tools=tools,
    verbose=True,
    memory=memory
)

from langchain_ibm import WatsonxLLM
 
# Параметры для модели
param = {
  "decoding_method": "greedy",
  "temperature": 0,
  "min_new_tokens": 5,
  "max_new_tokens": 250,
  "stop_sequences": ["\n\n"]
}
 
# Инициализация модели
model = WatsonxLLM(
  model_id = "ibm/granite-13b-chat-v2",
  url = credentials.get("url"),
  apikey = credentials.get("apikey"),
  project_id = credentials.get("project_id"),
  params = param
)

# Параметры для инструментального агента
instrument_params = {
  "decoding_method": "greedy",
  "temperature": 0,  # Минимальная вариативность
  "max_new_tokens": 300,
  "stop_sequences": ["\n\nObservation:"]  # Важно для ReAct паттерна!
}
 
# Параметры для креативного агента
creative_params = {
  "decoding_method": "sample",
  "temperature": 0.7,  # Больше вариативности
  "top_p": 0.9,
  "max_new_tokens": 500
}

precision_params = {
    "decoding_method": "greedy",
    "temperature": 0,
    "top_k": 1,
    "repetition_penalty": 1.0,
    "max_new_tokens": 200
}

creative_params = {
    "decoding_method": "sample",
    "temperature": 0.8,
    "top_p": 0.9,
    "top_k": 50,
    "repetition_penalty": 1.1,
    "max_new_tokens": 500
}

conversation_params = {
    "decoding_method": "sample",
    "temperature": 0.4,
    "top_p": 0.85,
    "repetition_penalty": 1.07,
    "max_new_tokens": 300
}

react_params = {
    "decoding_method": "greedy",
    "temperature": 0.1,  # Небольшая вариативность все же нужна
    "max_new_tokens": 250,
    "stop_sequences": ["Observation:", "Human:", "\n\n"]
}

code_params = {
    "decoding_method": "greedy",
    "temperature": 0.2,
    "repetition_penalty": 1.2,
    "max_new_tokens": 400
}

# Создаем виртуальное окружение
python -m venv langchain_env
 
# Активируем его
# На Windows:
langchain_env\Scripts\activate
# На Linux/Mac:
source langchain_env/bin/activate

WATSONX_API_KEY=ваш_ключ_api
WATSONX_URL=https://us-south.ml.cloud.ibm.com
WATSONX_PROJECT_ID=ваш_id_проекта
NASA_API_KEY=DEMO_KEY  # Или ваш собственный ключ

import langchain
import nasapy
print(f"LangChain версия: {langchain.__version__}")
n = nasapy.Nasa(key='DEMO_KEY')
print("Подключение к NASA API успешно!")

import os
from dotenv import load_dotenv
 
# Загружаем переменные из .env файла
load_dotenv()
 
# Теперь можно получить переменные через os.environ
api_key = os.environ.get("WATSONX_API_KEY")
project_id = os.environ.get("WATSONX_PROJECT_ID")
url = os.environ.get("WATSONX_URL")
 
# Проверяем, что ключи загрузились
if not api_key or not project_id:
    raise ValueError("Не удалось загрузить API ключи из .env файла")

import os
import json
from dotenv import load_dotenv
 
class Config:
    def __init__(self, config_file=None):
        # Загружаем переменные окружения
        load_dotenv()
        
        # Базовые настройки
        self.api_key = os.environ.get("WATSONX_API_KEY")
        self.project_id = os.environ.get("WATSONX_PROJECT_ID")
        self.url = os.environ.get("WATSONX_URL", "https://us-south.ml.cloud.ibm.com")
        
        # Загружаем дополнительные настройки из файла, если он указан
        if config_file and os.path.exists(config_file):
            with open(config_file, "r") as f:
                config_data = json.load(f)
                self.__dict__.update(config_data)
                
    def validate(self):
        """Проверяем, что все необходимые настройки присутствуют"""
        required = ["api_key", "project_id"]
        missing = [field for field in required if not getattr(self, field, None)]
        if missing:
            raise ValueError(f"Отсутствуют обязательные настройки: {', '.join(missing)}")
        return True

# Создаем и валидируем конфигурацию
config = Config("config.json")
config.validate()
 
# Используем настройки
model = WatsonxLLM(
    model_id="ibm/granite-13b-chat-v2",
    url=config.url,
    apikey=config.api_key,
    project_id=config.project_id,
    params=param
)

import os
from dotenv import load_dotenv
 
# Определяем текущее окружение
env = os.environ.get("ENVIRONMENT", "dev")
 
# Загружаем соответствующий файл
load_dotenv(f".env.{env}")

class ApiKeyManager:
    def __init__(self, keys_dict):
        self._keys = {service: keys for service, keys in keys_dict.items()}
        self._usage_count = {service: {key: 0 for key in keys} for service, keys in self._keys.items()}
    
    def get_key(self, service):
        """Возвращает ключ с наименьшим использованием"""
        if service not in self._keys:
            raise ValueError(f"Сервис {service} не найден")
        
        keys = self._keys[service]
        usage = self._usage_count[service]
        
        # Выбираем ключ с минимальным количеством использований
        key = min(keys, key=lambda k: usage[k])
        usage[key] += 1
        
        return key

import time
import random
from functools import wraps
 
def retry_with_backoff(max_retries=3, base_delay=1, max_delay=60):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            retries = 0
            while retries <= max_retries:
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    if "rate limit" not in str(e).lower() and "too many requests" not in str(e).lower():
                        raise  # Пробрасываем ошибку, если она не связана с rate limiting
                    
                    retries += 1
                    if retries > max_retries:
                        raise
                    
                    # Экспоненциальная задержка с случайным фактором
                    delay = min(base_delay * (2 ** (retries - 1)) + random.uniform(0, 1), max_delay)
                    print(f"Rate limit превышен. Повторная попытка через {delay:.2f} секунд...")
                    time.sleep(delay)
            return func(*args, **kwargs)
        return wrapper
    return decorator

@tool
@retry_with_backoff()
def get_astronomy_image(date: str):
    """Получает астрономическое изображение дня NASA на указанную дату."""
    apod = n.picture_of_the_day(date, hd=True)
    return apod['url']

@tool
def get_todays_date() -> str:
  """Получает текущую дату в формате ГГГГ-ММ-ДД."""
  from datetime import datetime
  date = datetime.now().strftime("%Y-%m-%d")
  return date
 
@tool(return_direct=True)
def get_astronomy_image(date: str):
  """Получает астрономическое изображение дня NASA на указанную дату. Дата в формате ГГГГ-ММ-ДД."""
  apod = n.picture_of_the_day(date, hd=True)
  return apod['url']
 
tools = [get_todays_date, get_astronomy_image]

system_prompt = """Отвечай пользователю максимально полезно и точно. У тебя есть доступ к следующим инструментам: {tools}
Используй JSON для вызова инструмента, указывая "action" и "action_input".
Допустимые значения для "action": "Final Answer" или {tool_names}
 
Формат должен быть таким:
{{
"action": НАЗВАНИЕ_ИНСТРУМЕНТА,
"action_input": ВХОДНЫЕ_ДАННЫЕ
}}
 
Следуй этому формату:
Вопрос: вопрос пользователя
Мысль: обдумай задачу
Действие:
$JSON_БЛОК
Наблюдение: результат действия
... (повтори Мысль/Действие/Наблюдение сколько нужно)
Мысль: я знаю ответ
Действие:
{{
"action": "Final Answer",
"action_input": "Итоговый ответ пользователю"
}}
"""

human_prompt = """{input}
{agent_scratchpad}"""
 
prompt = ChatPromptTemplate.from_messages([
  ("system", system_prompt),
  MessagesPlaceholder("chat_history", optional=True),
  ("human", human_prompt),
])
 
# Добавим информацию об инструментах
prompt = prompt.partial(
  tools=render_text_description_and_args(list(tools)),
  tool_names=", ".join([t.name for t in tools]),
)

memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

chain = (RunnablePassthrough.assign(
    agent_scratchpad=lambda x: format_log_to_str(x["intermediate_steps"]),
    chat_history=lambda x: memory.chat_memory.messages,
  )
  | prompt | model | JSONAgentOutputParser()
)
 
agent_executor = AgentExecutor(agent=chain, tools=tools, verbose=True, memory=memory)

response = agent_executor.invoke({"input": "Покажи мне астрономическую картинку дня от NASA на сегодня"})
print(response["output"])

@tool
def get_weather(location: str) -> str:
    """Получает текущую погоду для указанного местоположения."""
    api_key = os.getenv("WEATHER_API_KEY")
    url = f"https://api.weatherapi.com/v1/current.json?key={api_key}&q={location}&aqi=no"
    
    response = requests.get(url)
    if response.status_code != 200:
        return f"Ошибка при получении погоды: {response.status_code}"
    
    data = response.json()
    return f"Температура: {data['current']['temp_c']}°C, Условия: {data['current']['condition']['text']}"

from langchain.utilities import SQLDatabase
from langchain_community.tools import SQLDatabaseTool
 
# Подключаемся к базе данных
db = SQLDatabase.from_uri("postgresql://user:password@localhost:5432/mydatabase")
 
# Создаем инструмент для работы с БД
sql_tool = SQLDatabaseTool(db=db)
 
# Добавляем его в список инструментов
tools.append(sql_tool)

@tool
def read_file(file_path: str) -> str:
    """Читает содержимое файла по указанному пути."""
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            return f.read()
    except Exception as e:
        return f"Ошибка при чтении файла: {str(e)}"
 
@tool
def list_directory(directory_path: str) -> str:
    """Показывает содержимое директории."""
    try:
        files = os.listdir(directory_path)
        return "\n".join(files)
    except Exception as e:
        return f"Ошибка при чтении директории: {str(e)}"

from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 
# Инициализируем модель с поддержкой стриминга
streaming_model = WatsonxLLM(
  model_id="ibm/granite-13b-chat-v2",
  url=credentials.get("url"),
  apikey=credentials.get("apikey"),
  project_id=credentials.get("project_id"),
  params=param,
  streaming=True,  # Включаем стриминг!
  callbacks=[StreamingStdOutCallbackHandler()]
)
 
# Используем модель с поддержкой стриминга в цепочке
streaming_chain = (RunnablePassthrough.assign(
  agent_scratchpad=lambda x: format_log_to_str(x["intermediate_steps"]),
  chat_history=lambda x: memory.chat_memory.messages,
) | prompt | streaming_model | JSONAgentOutputParser())
 
streaming_agent = AgentExecutor(agent=streaming_chain, tools=tools, verbose=True)

def stream_agent_response(query):
  chunks = []
  
  def collector(chunk):
      chunks.append(chunk)
      yield chunk
  
  # Создаем кастомный обработчик для сбора чанков
  callback = StreamingCallbackHandler(collector)
  
  # Запускаем агента с этим обработчиком
  streaming_agent.invoke(
      {"input": query},
      config={"callbacks": [callback]}
  )
  
  # Возвращаем полный ответ в конце
  return "".join(chunks)

import asyncio
 
async def process_multiple_queries(queries):
  # Запускаем обработку всех запросов параллельно
  tasks = [streaming_agent.ainvoke({"input": q}) for q in queries]
  results = await asyncio.gather(*tasks)
  return [r["output"] for r in results]
 
# Использование
queries = ["Какая сегодня дата?", "Покажи астрономическую картинку за вчера"]
results = asyncio.run(process_multiple_queries(queries))

from langchain.vectorstores import Chroma
from langchain.embeddings import WatsonxEmbeddings
 
# Инициализируем модель для создания эмбеддингов
embeddings = WatsonxEmbeddings(
    url=credentials.get("url"),
    apikey=credentials.get("apikey"),
    project_id=credentials.get("project_id"),
    model_id="ibm/granite-embeddings"
)
 
# Создаем или подключаемся к векторному хранилищу
vectorstore = Chroma(
    collection_name="knowledge_base",
    embedding_function=embeddings,
    persist_directory="./chroma_db"
)
 
# Добавляем документы в базу
vectorstore.add_texts(["Земля — третья планета от Солнца", "Марс — четвертая планета от Солнца"])

@tool
def search_knowledge_base(query: str) -> str:
    """Ищет информацию в базе знаний по запросу."""
    results = vectorstore.similarity_search(query, k=2)
    if not results:
        return "Информация не найдена."
    return "\n".join([doc.page_content for doc in results])

from langchain.chains import RetrievalQA
 
# Создаем объект для поиска
retriever = vectorstore.as_retriever()
 
# Строим RAG-цепочку
rag_chain = RetrievalQA.from_chain_type(
    llm=model,
    chain_type="stuff",  # "stuff" - все найденные документы передаются в промпт
    retriever=retriever,
    verbose=True
)
 
# Используем цепочку
answer = rag_chain.invoke({"query": "Какая планета четвертая от Солнца?"})

def generate_search_queries(question):
    # Генерируем несколько поисковых запросов для исходного вопроса
    query_gen_prompt = f"""
    На основе вопроса пользователя, создай 3 поисковых запроса для поиска релевантной информации.
    Запросы должны быть краткими и содержать ключевые слова.
    
    Вопрос: {question}
    
    Поисковые запросы:"""
    
    response = model.invoke(query_gen_prompt)
    queries = [q.strip() for q in response.split("\n") if q.strip()]
    return queries[:3]  # Берем первые три запроса

def enhanced_rag_search(question):
    queries = generate_search_queries(question)
    all_docs = []
    
    for query in queries:
        docs = vectorstore.similarity_search(query, k=2)
        all_docs.extend(docs)
    
    # Удаляем дубликаты
    unique_docs = list({doc.page_content: doc for doc in all_docs}.values())
    
    # Возвращаем контекст
    return "\n".join([doc.page_content for doc in unique_docs])

def decompose_complex_query(query):
    decomposition_prompt = f"""
    Раздели следующий сложный запрос на последовательность простых подзадач:
    {query}
    
    Подзадачи:"""
    
    tasks = model.invoke(decomposition_prompt).strip().split("\n")
    return [task.strip() for task in tasks if task.strip()]

class ContextManager:
    def __init__(self):
        self.context = {}
        
    def set(self, key, value):
        self.context[key] = value
        
    def get(self, key, default=None):
        return self.context.get(key, default)
        
    def get_all_as_text(self):
        return "\n".join([f"{k}: {v}" for k, v in self.context.items()])

def self_reflection_enhancement(draft_answer, question):
    reflection_prompt = f"""
    Вопрос: {question}
    Черновик ответа: {draft_answer}
    
    Проанализируй черновик ответа:
    1. Отвечает ли он на все аспекты вопроса?
    2. Есть ли логические ошибки?
    3. Достаточно ли контекста?
    
    Улучшенный ответ:"""
    
    return model.invoke(reflection_prompt)

domain_specific_template = """Ты эксперт в области {domain}. 
Используй профессиональную терминологию из этой сферы: {terminology}.
При ответе опирайся на основные концепции: {key_concepts}.
Формат ответа должен соответствовать стандартам этой области.
 
Запрос пользователя: {query}
"""
 
legal_prompt = domain_specific_template.format(
    domain="юриспруденция",
    terminology="иск, ответчик, истец, арбитраж, юрисдикция",
    key_concepts="презумпция невиновности, состязательность сторон, бремя доказывания",
    query=user_query
)

medical_prompt = f"""Ты медицинский консультант. Отвечай строго по медицинским стандартам.
 
Пример 1:
Вопрос: Болит голова и температура 38.
Ответ: Ваши симптомы могут указывать на вирусную инфекцию или грипп. Рекомендую обильное питье, жаропонижающее и консультацию врача, если симптомы не улучшатся в течение 24 часов.
 
Ваш вопрос: {query}
"""

from langchain_community.document_loaders import PyPDFLoader
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_ibm.chat_models import WatsonxChat
from langchain_openai import OpenAIVisionChat  # Для работы с изображениями
 
# Модель для текста
text_model = WatsonxChat(model_id="ibm/granite-13b-chat-v2")
# Модель для обработки изображений
vision_model = OpenAIVisionChat()

@tool
def analyze_image(image_url: str) -> str:
    """Анализирует содержимое изображения и возвращает его описание."""
    response = vision_model.invoke([
        {"type": "text", "text": "Опиши детально, что изображено на картинке."},
        {"type": "image_url", "image_url": image_url}
    ])
    return response.content
 
@tool
def extract_text_from_pdf(pdf_url: str) -> str:
    """Извлекает текст из PDF документа по указанной ссылке."""
    loader = PyPDFLoader(pdf_url)
    pages = loader.load()
    return "\n".join(page.page_content for page in pages)

def route_by_content_type(query):
    if "изображени" in query.lower() or "картинк" in query.lower() or "фото" in query.lower():
        return "vision"
    elif "pdf" in query.lower() or "документ" in query.lower() or "файл" in query.lower():
        return "document"
    else:
        return "text"
 
# А дальше уже используем это в цепочке принятия решений

def analyze_query_needs(query):
    analysis_prompt = f"""
    Определи, какие типы данных нужны для ответа на этот запрос:
    "{query}"
    
    Укажи необходимые типы (текст/изображения/документы):
    """
    
    analysis = text_model.invoke(analysis_prompt)
    return analysis

from langchain.memory import ConversationBufferWindowMemory, ConversationSummaryMemory
from langchain.memory.chat_message_histories import RedisChatMessageHistory
 
# Память с ограниченным окном (хранит только N последних сообщений)
window_memory = ConversationBufferWindowMemory(
  k=5,  # Хранить только 5 последних взаимодействий
  return_messages=True
)
 
# Память с автоматическим суммированием (экономит токены)
summary_memory = ConversationSummaryMemory(
  llm=model,  # Модель для создания саммари
  return_messages=True
)

message_history = RedisChatMessageHistory(
  url="redis://localhost:6379/0",
  session_id=user_id,  # Уникальный ID пользователя
  ttl=604800  # Время жизни в секундах (неделя)
)
 
persistent_memory = ConversationBufferMemory(
  memory_key="chat_history",
  chat_memory=message_history,
  return_messages=True
)

class PersistentAgentState:
  def __init__(self, user_id, storage_path="./states"):
      self.user_id = user_id
      self.storage_path = storage_path
      os.makedirs(storage_path, exist_ok=True)
      self.state_file = os.path.join(storage_path, f"{user_id}.json")
      self.state = self._load_state()
  
  def _load_state(self):
      if os.path.exists(self.state_file):
          with open(self.state_file, "r") as f:
              return json.load(f)
      return {}  # Пустое состояние по умолчанию
  
  def save(self):
      with open(self.state_file, "w") as f:
          json.dump(self.state, f)
  
  def set(self, key, value):
      self.state[key] = value
      self.save()
  
  def get(self, key, default=None):
      return self.state.get(key, default)

from celery import Celery
import os
 
# Настраиваем Celery с Redis в качестве брокера сообщений
app = Celery('langchain_tasks', 
             broker='redis://localhost:6379/1',
             backend='redis://localhost:6379/2')
 
# Определяем задачу для выполнения запроса к агенту
@app.task
def process_agent_query(query, user_id):
    # Загружаем или создаем агента
    agent = load_agent_for_user(user_id)
    
    # Выполняем запрос
    try:
        result = agent.invoke({"input": query})
        return {"status": "success", "output": result["output"]}
    except Exception as e:
        return {"status": "error", "message": str(e)}

# Запускаем задачу асинхронно
task = process_agent_query.delay("Расскажи о звездах", user_id="user123")
 
# Возвращаем ID задачи пользователю
return {"task_id": task.id, "message": "Запрос обрабатывается, проверьте результат через несколько секунд"}

def check_task_status(task_id):
    task = process_agent_query.AsyncResult(task_id)
    if task.ready():
        return task.result
    else:
        return {"status": "processing"}

def notify_task_completion(task_id, websocket_id):
    # Периодически проверяем статус задачи
    task = process_agent_query.AsyncResult(task_id)
    if task.ready():
        # Отправляем результат через вебсокет
        send_to_websocket(websocket_id, task.result)
        return True
    return False

@app.task
def update_agent_knowledge_base():
    # Код для обновления базы знаний
    pass
 
# Запускаем задачу каждые 24 часа
app.conf.beat_schedule = {
    'update-knowledge-daily': {
        'task': 'tasks.update_agent_knowledge_base',
        'schedule': 86400.0,  # 24 часа в секундах
    },
}

from langchain.graphs import Neo4jGraph
from langchain_community.graphs.graph_document import (
    Node, Relationship, GraphDocument
)
 
# Подключаемся к Neo4j
graph = Neo4jGraph(
    url="bolt://localhost:7687", 
    username="neo4j", 
    password="password"
)
 
# Создаем узлы и связи на основе извлеченной информации
def update_knowledge_graph(information, user_id):
    # Извлекаем сущности и отношения из текста
    extraction_prompt = f"""
    Извлеки из текста все упомянутые сущности (люди, места, концепции) 
    и отношения между ними:
    {information}
    """
    
    extracted = model.invoke(extraction_prompt)
    
    # Парсим результат (в реальном коде использовал бы структурированный формат)
    # Создаем документ для графа
    doc = GraphDocument(
        nodes=[
            Node(id="user123", type="User", properties={"name": "Алексей"}),
            Node(id="pref1", type="Preference", properties={"topic": "Астрономия"})
        ],
        relationships=[
            Relationship(source="user123", target="pref1", type="INTERESTED_IN")
        ]
    )
    
    # Добавляем в граф
    graph.add_graph_documents([doc])

@tool
def query_memory_graph(query: str) -> str:
    """Ищет информацию в графе знаний агента."""
    cypher_template = """
    MATCH (u:User {{id: '{user_id}'}})-[r]-(n)
    WHERE n.topic CONTAINS '{query}' OR n.content CONTAINS '{query}'
    RETURN n, r
    LIMIT 5
    """
    
    results = graph.query(cypher_template.format(user_id=current_user_id, query=query))
    return format_graph_results(results)

from prometheus_client import Counter, Histogram, start_http_server
 
# Метрики для LangChain агента
llm_requests_total = Counter('llm_requests_total', 'Total number of LLM requests', ['model', 'status'])
llm_request_duration = Histogram('llm_request_duration_seconds', 'Duration of LLM requests', ['model'])
tool_calls_total = Counter('tool_calls_total', 'Total number of tool calls', ['tool_name', 'status'])
 
# Запускаем сервер метрик
start_http_server(8001)

from langchain.globals import set_llm_cache
from langchain.cache import InMemoryCache, RedisCache
 
# Для локальной разработки
set_llm_cache(InMemoryCache())
 
# Для продакшена с Redis
set_llm_cache(RedisCache(redis_url="redis://localhost:6379"))

class TieredCache:
  def __init__(self):
      self.memory_cache = {}  # Самый быстрый уровень
      self.redis_client = redis.Redis()  # Второй уровень
      
  def get(self, key):
      # Сначала проверяем in-memory кэш
      if key in self.memory_cache:
          return self.memory_cache[key]
      
      # Затем Redis
      value = self.redis_client.get(key)
      if value:
          # Обновляем in-memory кэш
          self.memory_cache[key] = value
          return value
      
      return None

def create_optimized_prompt(query, tools_only_needed):
  # Выбираем только нужные инструменты для конкретного запроса
  relevant_tools = [t for t in all_tools if is_tool_relevant(t, query)]
  
  # Генерируем компактное описание инструментов
  tools_description = "\n".join([f"- {t.name}: {t.description.split('.')[0]}" for t in relevant_tools])
  
  # Создаем оптимизированный промпт
  return f"Ты помощник. Инструменты: {tools_description}. Запрос: {query}"

async def batch_llm_requests(prompts, max_batch_size=5):
  results = []
  
  for i in range(0, len(prompts), max_batch_size):
      batch = prompts[i:i+max_batch_size]
      tasks = [model.ainvoke(prompt) for prompt in batch]
      batch_results = await asyncio.gather(*tasks)
      results.extend(batch_results)
  
  return results

def smart_retry(func, max_retries=3, base_delay=1, timeout=10):
  async def wrapped(*args, **kwargs):
      retries = 0
      last_exception = None
      
      while retries <= max_retries:
          try:
              return await asyncio.wait_for(func(*args, **kwargs), timeout=timeout)
          except asyncio.TimeoutError:
              last_exception = TimeoutError("Запрос превысил тайм-аут")
          except Exception as e:
              last_exception = e
          
          retries += 1
          delay = base_delay * (2 ** (retries - 1))
          await asyncio.sleep(delay)
      
      raise last_exception
  
  return wrapped

def verify_llm_response(response, query):
    """Проверяет ответ LLM на наличие галлюцинаций."""
    # Извлекаем ключевые утверждения из ответа
    verification_prompt = f"""
    Извлеки из ответа все фактические утверждения, которые можно проверить:
    {response}
    
    Формат: список утверждений, по одному на строку
    """
    claims = model.invoke(verification_prompt).strip().split("\n")
    
    results = []
    for claim in claims:
        # Проверяем каждое утверждение через поиск или другие инструменты
        search_results = search_tool.invoke(claim)
        confidence = calculate_confidence(claim, search_results)
        results.append({"claim": claim, "confidence": confidence})
    
    # Если есть утверждения с низкой уверенностью, возможно это галлюцинация
    suspicious = [r for r in results if r["confidence"] < 0.3]
    return suspicious

def self_evaluation_check(response, query):
    eval_prompt = f"""
    Вопрос: {query}
    Твой ответ: {response}
    
    Оцени достоверность своего ответа по шкале от 0 до 10.
    Укажи части ответа, в которых ты не уверен.
    Приведи рассуждение, почему ты считаешь ответ достоверным или недостоверным.
    """
    
    evaluation = model.invoke(eval_prompt)
    
    # Извлекаем числовую оценку из ответа
    confidence_match = re.search(r'(\d+)(?:/10)?', evaluation)
    if confidence_match:
        confidence = int(confidence_match.group(1)) / 10
        if confidence < 0.6:
            return {"suspicious": True, "confidence": confidence, "reasoning": evaluation}
    
    return {"suspicious": False, "evaluation": evaluation}

def uncertainty_detection(response):
    uncertainty_markers = [
        "возможно", "вероятно", "предположительно", "кажется",
        "я думаю", "я полагаю", "насколько я знаю", "я не уверен"
    ]
    
    score = sum([1 for marker in uncertainty_markers if marker in response.lower()])
    return score > 2  # Если больше двух маркеров, это подозрительно

def test_agent_response(query, agent, validators=None):
    """Тестирует ответ агента с применением различных валидаторов."""
    if validators is None:
        validators = [basic_format_validator, factual_consistency_validator]
    
    response = agent.invoke({"input": query})
    
    results = {}
    for validator in validators:
        validator_name = validator.__name__
        try:
            results[validator_name] = validator(query, response)
        except Exception as e:
            results[validator_name] = {"status": "error", "message": str(e)}
    
    return {
        "query": query,
        "response": response,
        "validation_results": results,
        "passed": all(r.get("status") == "pass" for r in results.values() if isinstance(r, dict))
    }

def basic_format_validator(query, response):
    """Проверяет базовый формат ответа."""
    if not response or not response.get("output"):
        return {"status": "fail", "message": "Пустой ответ"}
    
    output = response["output"]
    
    # Проверка на минимальную длину
    if len(output) < 10:
        return {"status": "fail", "message": "Ответ слишком короткий"}
    
    # Проверка на наличие типичных ошибок
    error_patterns = [
        r"я не могу|не имею доступа|не знаю|не уверен",
        r"ошибка|error|exception|traceback",
        r"action_input.*?None"
    ]
    
    for pattern in error_patterns:
        if re.search(pattern, output, re.IGNORECASE):
            return {
                "status": "warn", 
                "message": f"Обнаружен паттерн ошибки: {pattern}"
            }
    
    return {"status": "pass"}

def factual_consistency_validator(query, response):
    """Проверяет фактическую точность ответа."""
    output = response["output"]
    
    # Извлекаем утверждения из ответа
    verification_prompt = f"""
    Извлеки из ответа все фактические утверждения:
    {output}
    
    Выведи только список утверждений, по одному на строку.
    """
    
    try:
        statements = model.invoke(verification_prompt).strip().split('\n')
        if not statements or statements[0] == '':
            return {"status": "pass", "message": "Нет проверяемых утверждений"}
        
        # Проверяем каждое утверждение
        suspicious = []
        for statement in statements:
            if statement and len(statement) > 10:
                verification = verify_statement(statement)
                if verification["score"] < 0.3:
                    suspicious.append({
                        "statement": statement,
                        "score": verification["score"]
                    })
        
        if suspicious:
            return {
                "status": "warn",
                "message": "Обнаружены подозрительные утверждения",
                "details": suspicious
            }
        
        return {"status": "pass"}
    except Exception as e:
        return {"status": "error", "message": str(e)}

import logging
 
# Настраиваем логгер
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("agent_debug.log"),
        logging.StreamHandler()
    ]
)
 
# Создаем кастомный колбэк для LangChain
class DetailedLogger(BaseCallbackHandler):
    def on_llm_start(self, serialized, prompts, **kwargs):
        logging.debug(f"Запрос к LLM: {prompts[0][:100]}...")
        
    def on_llm_end(self, response, **kwargs):
        logging.debug(f"Ответ LLM: {response.generations[0][0].text[:100]}...")
        
    def on_tool_start(self, serialized, input_str, **kwargs):
        logging.info(f"Вызов инструмента {serialized['name']}: {input_str}")
        
    def on_tool_end(self, output, **kwargs):
        logging.info(f"Результат инструмента: {output}")
        
    def on_chain_start(self, serialized, inputs, **kwargs):
        logging.debug(f"Начало цепочки: {serialized['name']}")
        
    def on_chain_end(self, outputs, **kwargs):
        logging.debug(f"Конец цепочки: {outputs}")

# Инициализируем агента с логгером
agent_executor = AgentExecutor(
    agent=chain,
    tools=tools,
    verbose=True,
    callbacks=[DetailedLogger()]
)

def visualize_agent_reasoning(steps):
    from graphviz import Digraph
    
    dot = Digraph(comment='Ход рассуждений агента')
    
    for i, step in enumerate(steps):
        dot.node(f'thought_{i}', f'Мысль: {step["thought"]}')
        dot.node(f'action_{i}', f'Действие: {step["action"]}')
        dot.node(f'result_{i}', f'Результат: {step["observation"]}')
        
        if i > 0:
            dot.edge(f'result_{i-1}', f'thought_{i}')
        dot.edge(f'thought_{i}', f'action_{i}')
        dot.edge(f'action_{i}', f'result_{i}')
    
    return dot

# Слишком маленькое значение
bad_params = {
  "max_new_tokens": 50  # Недостаточно для сложных ответов
}
 
# Более адекватная настройка
better_params = {
  "max_new_tokens": 250,  # Больше пространства
  "stop_sequences": ["Observation:", "Human:"]  # Явные маркеры окончания
}

def truncate_safely(text, max_tokens=1000):
  """Безопасно обрезает текст до указанного количества токенов."""
  encoding = tiktoken.get_encoding("cl100k_base")  # Подходящий энкодер
  tokens = encoding.encode(text)
  
  if len(tokens) <= max_tokens:
      return text
      
  # Оставляем место для [обрезано] и важного начала/конца
  keep_start = int(max_tokens * 0.8)
  keep_end = max_tokens - keep_start - 10
  
  truncated = encoding.decode(tokens[:keep_start]) + "\n[...сокращено...]\n" + encoding.decode(tokens[-keep_end:])
  return truncated

import os
import nasapy
from datetime import datetime
from dotenv import load_dotenv
 
from langchain_core.tools import tool
from langchain.agents import AgentExecutor
from langchain.agents.format_scratchpad import format_log_to_str
from langchain.agents.output_parsers import JSONAgentOutputParser
from langchain.memory import ConversationBufferMemory
from langchain.tools.render import render_text_description_and_args
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from langchain_ibm import WatsonxLLM
 
# Загружаем переменные окружения
load_dotenv()
 
# Инициализируем NASA API
nasa = nasapy.Nasa(key=os.getenv("NASA_API_KEY", "DEMO_KEY"))
 
# Определяем инструменты
@tool
def get_todays_date() -> str:
    """Получает текущую дату в формате ГГГГ-ММ-ДД."""
    date = datetime.now().strftime("%Y-%m-%d")
    return date
 
@tool(return_direct=True)
def get_astronomy_image(date: str):
    """Получает астрономическое изображение дня NASA на указанную дату. Дата в формате ГГГГ-ММ-ДД."""
    apod = nasa.picture_of_the_day(date, hd=True)
    return apod['url']
 
tools = [get_todays_date, get_astronomy_image]