Final_GAIA_test / tmp.txt
DrZimmer's picture
Upload 10 files
568dd3c verified
raw
history blame
54.2 kB
import json
import os
import pandas as pd
import PyPDF2
import requests
from PIL import Image
from pathlib import Path
from langgraph.graph import StateGraph, END
from typing import Dict, Any, TypedDict, Optional
from docx import Document
from pptx import Presentation
from langchain_ollama import ChatOllama
import logging
import importlib.util
import re
import pydub
import xml.etree.ElementTree as ET
from concurrent.futures import ThreadPoolExecutor, TimeoutError
from duckduckgo_search import DDGS
from tqdm import tqdm
import pytesseract
import torch
from faster_whisper import WhisperModel
from sentence_transformers import SentenceTransformer
import faiss
from faiss import IndexFlatL2
import ollama
import asyncio
from shazamio import Shazam
from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
from bs4 import BeautifulSoup
from retrying import retry
import pdfplumber
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
# Настройка путей для Hugging Face Spaces
BASE_DIR = "/home/user/app" # Базовая директория в Hugging Face Spaces
# --- Константы ---
DATA_DIR = os.path.join(BASE_DIR, "2023")
TEMP_DIR = os.path.join(BASE_DIR, "temp")
# Константы
METADATA_PATH = os.path.join(BASE_DIR, "metadata.jsonl")
OLLAMA_URL = "http://localhost:11434" # Ollama в контейнере
MODEL_NAME = "qwen2:7b"
ANSWERS_PATH = os.path.join(BASE_DIR, "answers.json")
UNKNOWN_PATH = os.path.join(BASE_DIR, "unknown.txt")
TRANSCRIPTION_TIMEOUT = 30
MAX_AUDIO_DURATION = 300
ANSWERS_JSON = "answers.json"
UNKNOWN_FILE = "unknown.txt"
# Создание временной папки
if not os.path.exists(TEMP_DIR):
os.makedirs(TEMP_DIR)
# Настройка Tesseract
pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract" # Путь в контейнере
# Настройка логгирования
LOG_FILE = os.path.join(BASE_DIR, "log.txt")
logging.basicConfig(
filename=LOG_FILE,
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
filemode="w"
)
logger = logging.getLogger(__name__)
# Отключаем отладочные логи от сторонних библиотек
logging.getLogger("sentence_transformers").setLevel(logging.WARNING)
logging.getLogger("faster_whisper").setLevel(logging.WARNING)
logging.getLogger("faiss").setLevel(logging.WARNING)
logging.getLogger("ctranslate2").setLevel(logging.WARNING)
logging.getLogger("torch").setLevel(logging.WARNING)
logging.getLogger("pydub").setLevel(logging.WARNING)
logging.getLogger("shazamio").setLevel(logging.WARNING)
# --- Проверка зависимостей ---
def check_openpyxl():
if importlib.util.find_spec("openpyxl") is None:
logger.error("openpyxl не установлена. Установите: pip install openpyxl")
raise ImportError("openpyxl не установлена. Установите: pip install openpyxl")
logger.info("openpyxl доступна.")
def check_pydub():
if importlib.util.find_spec("pydub") is None:
logger.error("pydub не установлена. Установите: pip install pydub")
raise ImportError("pydub не установлена. Установите: pip install pydub")
logger.info("pydub доступна.")
def check_faster_whisper():
if importlib.util.find_spec("faster_whisper") is None:
logger.error("faster-whisper не установлена. Установите: pip install faster-whisper")
raise ImportError("faster-whisper не установлена. Установите: pip install faster-whisper")
logger.info("faster-whisper доступна.")
def check_sentence_transformers():
if importlib.util.find_spec("sentence_transformers") is None:
logger.error("sentence-transformers не установлена. Установите: pip install sentence-transformers")
raise ImportError("sentence-transformers не установлена. Установите: pip install sentence-transformers")
logger.info("sentence-transformers доступна.")
def check_faiss():
if importlib.util.find_spec("faiss") is None:
logger.error("faiss не установлена. Установите: pip install faiss-cpu")
raise ImportError("faiss не установлена. Установите: pip install faiss-cpu")
logger.info("faiss доступна.")
def check_ollama():
if importlib.util.find_spec("ollama") is None:
logger.error("ollama не установлена. Установите: pip install ollama")
raise ImportError("ollama не установлена. Установите: pip install ollama")
logger.info("ollama доступна.")
def check_shazamio():
if importlib.util.find_spec("shazamio") is None:
logger.error("shazamio не установлена. Установите: pip install shazamio")
raise ImportError("shazamio не установлена. Установите: pip install shazamio")
logger.info("shazamio доступна.")
def check_langchain_community():
if importlib.util.find_spec("langchain_community") is None:
logger.error("langchain_community не установлена. Установите: pip install langchain-community")
raise ImportError("langchain_community не установлена. Установите: pip install langchain-community")
logger.info("langchain_community доступна.")
# Инициализация модели
try:
llm = ChatOllama(base_url=OLLAMA_URL, model=MODEL_NAME, request_timeout=60)
test_response = llm.invoke("Test")
if test_response is None or not hasattr(test_response, 'content'):
raise ValueError("Ollama модель недоступна или возвращает некорректный ответ")
logger.info("Модель ChatOllama инициализирована.")
except Exception as e:
logger.error(f"Ошибка инициализации модели: {e}")
raise e
# --- Состояние для LangGraph ---
class AgentState(TypedDict):
question: str
task_id: str
file_path: Optional[str]
file_content: Optional[str]
wiki_results: Optional[str]
arxiv_results: Optional[str]
web_results: Optional[str]
answer: str
raw_answer: str
# --- Функция извлечения тайминга ---
def extract_timing(question: str) -> int:
"""
Извлекает тайминг (в миллисекундах) из вопроса.
Поддерживает форматы: '2-minute', '2 minutes', '2 min mark', '120 seconds', '1 min 30 sec'.
Если тайминг не найден, возвращает 0 (обрезка с начала на 20 секунд).
"""
question = question.lower()
total_ms = 0
# Поиск минут (2-minute, 2 minutes, 2 min, 2 min mark, etc.)
minute_match = re.search(r'(\d+)\s*(?:-|\s)?\s*(?:minute|min)\b(?:\s*mark)?', question)
if minute_match:
minutes = int(minute_match.group(1))
total_ms += minutes * 60 * 1000
# Поиск секунд (120 seconds, 30 sec, etc.)
second_match = re.search(r'(\d+)\s*(?:second|sec|s)\b', question)
if second_match:
seconds = int(second_match.group(1))
total_ms += seconds * 1000
logger.info(f"Extracted timing: {total_ms // 60000} minutes, {(total_ms % 60000) // 1000} seconds ({total_ms} ms)")
return total_ms
# --- Функция распознавания песни ---
async def recognize_song(audio_file: str, start_time_ms: int = 0, duration_ms: int = 20000) -> dict:
try:
logger.info(f"Trimming audio from {start_time_ms/1000:.2f} seconds...")
audio = pydub.AudioSegment.from_file(audio_file, format="mp3")
end_time_ms = start_time_ms + duration_ms
if end_time_ms > len(audio):
end_time_ms = len(audio)
trimmed_audio = audio[start_time_ms:end_time_ms]
trimmed_path = os.path.join(TEMP_DIR, "trimmed_song.wav")
trimmed_audio.export(trimmed_path, format="wav")
logger.info(f"Trimmed audio saved to {trimmed_path}")
logger.info("Recognizing song with Shazam...")
shazam = Shazam()
result = await shazam.recognize_song(trimmed_path)
track = result.get("track", {})
title = track.get("title", "Not found")
artist = track.get("subtitle", "Unknown")
logger.info(f"Shazam result: Title: {title}, Artist: {artist}")
return {"title": title, "artist": artist}
except Exception as e:
logger.error(f"Error recognizing song: {str(e)}")
return {"title": "Not found", "artist": "Unknown"}
# --- Функция транскрипции MP3 ---
def transcribe_audio(audio_file: str, chunk_length_ms: int = 300000) -> str:
"""
Транскрибирует MP3-файл и возвращает полный текст.
Args:
audio_file: Путь к MP3-файлу.
chunk_length_ms: Длина чанка в миллисекундах (по умолчанию 300000, т.е. 5 минут).
Returns:
Полный текст или сообщение об ошибке.
"""
logger.info(f"Начало транскрипции файла: {audio_file}")
try:
if not os.path.exists(audio_file):
logger.error(f"Файл {audio_file} не найден")
return f"Error: Audio file {audio_file} not found in {os.getcwd()}"
logger.info(f"Инициализация WhisperModel для {audio_file}")
device = "cuda" if torch.cuda.is_available() else "cpu"
model = WhisperModel("small", device=device, compute_type="float16" if device == "cuda" else "int8")
logger.info("Модель Whisper инициализирована")
logger.info(f"Загрузка аудио: {audio_file}")
audio = pydub.AudioSegment.from_file(audio_file)
logger.info(f"Длительность аудио: {len(audio)/1000:.2f} секунд")
chunks = []
temp_dir = os.path.join(TEMP_DIR, "audio_chunks")
os.makedirs(temp_dir, exist_ok=True)
logger.info(f"Создана временная папка: {temp_dir}")
for i in range(0, len(audio), chunk_length_ms):
chunk = audio[i:i + chunk_length_ms]
chunk_file = os.path.join(temp_dir, f"chunk_{i//chunk_length_ms}.mp3")
chunk.export(chunk_file, format="mp3")
chunks.append(chunk_file)
logger.info(f"Создан чанк {i+1}: {chunk_file}")
logger.info(f"Создано {len(chunks)} чанков")
full_text = []
chunks_text = []
for i, chunk in enumerate(tqdm(chunks, desc="Transcribing chunks")):
logger.info(f"Обработка чанка {i+1}/{len(chunks)}: {chunk}")
segments, _ = model.transcribe(chunk, language="en")
chunk_text = " ".join(segment.text for segment in segments).strip()
full_text.append(chunk_text)
chunks_text.append(f"Chunk-{i+1}:\n{chunk_text}\n---\n")
logger.info(f"Чанк {i+1} транскрибирован: {chunk_text[:50]}...")
logger.info("Транскрипция чанков завершена")
logger.info("Запись результатов транскрипции")
with open(os.path.join(TEMP_DIR, "chunks.txt"), "w", encoding="utf-8") as f:
f.write("\n".join(chunks_text))
combined_text = " ".join(full_text)
with open(os.path.join(TEMP_DIR, "total_text.txt"), "w", encoding="utf-8") as f:
f.write(combined_text)
logger.info("Результаты транскрипции записаны")
word_count = len(combined_text.split())
token_count = int(word_count * 1.3)
logger.info(f"Транскрибировано: {word_count} слов, ~{token_count} токенов")
logger.info("Очистка временных файлов")
for chunk_file in chunks:
if os.path.exists(chunk_file):
os.remove(chunk_file)
logger.info(f"Удален чанк: {chunk_file}")
if os.path.exists(temp_dir):
os.rmdir(temp_dir)
logger.info(f"Удалена папка: {temp_dir}")
logger.info(f"Транскрипция завершена успешно: {audio_file}")
return combined_text
except Exception as e:
logger.error(f"Ошибка транскрипции аудио: {str(e)}")
return f"Error processing audio: {str(e)}"
# --- Создание RAG-индекса ---
def create_rag_index(text: str, model: SentenceTransformer) -> tuple:
sentences = [s.strip()[:500] for s in text.split(".") if s.strip()]
embeddings = model.encode(sentences, convert_to_numpy=True, show_progress_bar=False)
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)
return index, sentences, embeddings
# --- Обработка файлов ---
async def process_file(file_path: str, question: str) -> str:
if not file_path:
logger.warning("Файл не указан")
return "Файл не указан."
full_path = os.path.join(BASE_DIR, file_path) if file_path else None
if not full_path or not Path(full_path).exists():
logger.warning(f"Файл не найден: {full_path or file_path}")
return f"Файл не найден: {file_path}"
ext = Path(full_path).suffix.lower()
logger.info(f"Обработка файла: {full_path} (формат: {ext})")
try:
if ext == ".pdf":
try:
import pdfplumber
with pdfplumber.open(full_path) as pdf:
text = "".join(page.extract_text() or "" for page in pdf.pages)
if not text.strip():
logger.warning(f"Пустой текст в PDF: {full_path}")
return "Пустой PDF-файл"
return text
except ImportError:
logger.warning("pdfplumber не установлен. Используется PyPDF2.")
with open(full_path, "rb") as f:
reader = PyPDF2.PdfReader(f)
text = "".join(page.extract_text() or "" for page in reader.pages)
if not text.strip():
logger.warning(f"Пустой текст в PDF: {full_path}")
return "Пустой PDF-файл"
return text
elif ext in [".xlsx", ".csv"]:
if ext == ".xlsx":
check_openpyxl()
df = pd.read_excel(full_path) if ext == ".xlsx" else pd.read_csv(full_path)
if df.empty:
logger.warning(f"Пустой DataFrame для файла {full_path}")
return "Пустой файл"
return str(df.to_string())
elif ext in [".txt", ".json", ".jsonl"]:
with open(full_path, "r", encoding="utf-8") as f:
text = f.read()
if "how many" in question.lower():
numbers = re.findall(r'\b\d+\b', text)
if numbers:
logger.info(f"Найдены числа в тексте: {numbers}")
return f"Числа: {', '.join(numbers)}\nТекст: {text[:1000]}"
return text
elif ext in [".png", ".jpg"]:
try:
image = Image.open(full_path)
text = pytesseract.image_to_string(image)
if not text.strip():
logger.warning(f"Пустой текст в изображении: {full_path}")
return f"Изображение: {full_path} (OCR не дал результата)"
logger.info(f"OCR выполнен: {text[:50]}...")
return f"OCR текст: {text}"
except Exception as e:
logger.error(f"Ошибка OCR для {full_path}: {e}")
return f"Ошибка: {str(e)}"
elif ext == ".docx":
doc = Document(full_path)
return "\n".join(paragraph.text for paragraph in doc.paragraphs)
elif ext == ".pptx":
prs = Presentation(full_path)
text = ""
for slide in prs.slides:
for shape in slide.shapes:
if hasattr(shape, "text"):
text += shape.text + "\n"
return text
elif ext == ".mp3":
if "name of the song" in question.lower() or "what song" in question.lower():
check_shazamio()
check_pydub()
start_time_ms = extract_timing(question)
if start_time_ms == 0 and not re.search(r"(?:minute|min|second|sec|s)\b", question):
logger.info("No timing specified, using default 0–20 seconds")
result = await recognize_song(full_path, start_time_ms)
title = result["title"]
logger.info(f"Song recognition result: {title}")
return title
if "duration" in question.lower() or "minute" in question.lower():
try:
audio = pydub.audioSegment.audio_file(full_path)
duration = len(audio) // 1000
logger.info(f"Audio duration: {duration:.2f']} seconds")
return f"Duration: {duration:.2f} seconds"
except Exception as e:
logger.error(f"Error getting duration: {e}")
return f"Error: {e}"
except Exception as e:
logger.error(f"Ошибка получения длительности: {e}")
return f"Ошибка: {str(e)}"
check_faster_hhisper()
check_ccheerwer():
check_kick_faiss():
check_shick_ollama()
transcribed_text = transcribe_audio(full_path)
if transcribed_text.startswith("Error"):
logger.error(f"Ошибка транскрипции: {transcribed_text}")
return transcribed_text
return transcribed_text
elif ext == ".m4a":
if "how long" in question.lower() or "minute" in question.lower():
try:
audio = pydub.AudioSegment.from_file(full_path)
duration = len(audio) / 1000
logger.info(f"Длительность аудио: {duration:.2f} секунд")
return f"Длительность: {duration:.2f} секунд"
except Exception as e:
logger.error(f"Ошибка обработки: {e}")
return f"Ошибка: {str(e)}"
logger.warning(f"Транскрипция M4A не поддерживается для {full_path}")
return f"Аудиофайл: {full_path} (транскрипция не выполнена)"
elif ext == ".xml":
tree = ET.parse(full_path)
root = tree.getroot()
text = " ".join(elem.text or "" for elem in root.iter())
return text
else:
logger.warning(f"Формат не поддерживается: {ext}")
return f"Формат {ext} не поддерживается."
except Exception as e:
logger.error(f"Ошибка обработки файла {full_path}: {e}")
return f"Ошибка обработки файла: {str(e)}"
# --- Разбор текста PDF ---
def process_pdf(file_path: str) -> str:
"""Извлечение текста из PDF файла."""
try:
with pdfplumber.open(file_path) as pdf:
text = ""
for page in pdf.pages:
page_text = page.extract_text()
if page_text:
text += page_text + "\n"
return text.strip() if text else "No text extracted from PDF"
except Exception as e:
logger.error(f"Ошибка извлечения текста из PDF {file_path}: {str(e)}")
return f"Error extracting text from PDF: {str(e)}"
# --- Узлы LangGraph ---
def analyze_question(state: AgentState) -> AgentState:
logger.info(f"Вход в analyze_question, state: {state}")
if not isinstance(state, dict):
logger.error(f"analyze_question: state не является словарем: {state}")
return {"answer": "Error: Invalid state in analyze_question", "raw_answer": "Error: Invalid state in analyze_question"}
task_id = state.get("task_id", "unknown")
question = state.get("question", "")
file_path = state.get("file_path")
logger.info(f"Анализ задачи {task_id}: Вопрос: {question[:50]}...")
if file_path:
state["file_content"] = process_file(file_path, question)
else:
state["file_content"] = None
logger.info("Файл не указан для задачи.")
logger.info(f"Содержимое файла: {state['file_content'][:50] if state['file_content'] else 'Нет файла'}...")
logger.info(f"Выход из analyze_question, state: {state}")
return state
# --- Для US Census, Macrotrends, Twitter, музеев ---
def scrape_website(url, query):
"""Скрейпинг веб-сайта с повторными попытками."""
try:
headers = {"User-Agent": "Mozilla/5.0"}
response = requests.get(url, params={"q": query}, headers=headers, timeout=10)
soup = BeautifulSoup(response.text, "html.parser")
text = soup.get_text(separator=" ", strip=True)
return text[:1000] if text and len(text.strip()) > 50 else "No relevant content found"
except Exception as e:
logger.error(f"Ошибка парсинга {url}: {str(e)}")
return f"Error: {str(e)}"
# --- web поиск по категориям ---
def web_search(state: AgentState) -> AgentState:
logger.info(f"Вход в web_search, state: {state}")
if not isinstance(state, dict):
logger.error(f"web_search: state не является словарем: {type(state)}")
return {"answer": "Error: Invalid state in web_search", "raw_answer": "Error: Invalid state in web_search"}
question = state.get("question", "")
task_id = state.get("task_id", "unknown")
question_lower = question.lower()
logger.info(f"Поиск для задачи {task_id} в веб-поиске...")
try:
logger.info("Проверка доступности langchain_community...")
try:
from langchain_community.utils import WikipediaAPIWrapper, ArxivAPIWrapper
except ImportError as e:
logger.error(f"langchain_community не установлен: {str(e)}")
raise ImportError(f"langchain_community is not available: {str(e)}")
query = question[:500]
logger.info(f"Выполнение поиска для запроса: {query[:50]}...")
state["wiki_results"] = state.get("wiki_results", "")
state["arxiv_results"] = state.get("arxiv_results", "")
state["web_results"] = state.get("web_results", "")
state["file_content"] = state.get("file_content", "")
if "census" in question_lower:
logger.info("Поиск на US Census...")
content = scrape_website("https://www.census.gov", query)
state["web_results"] = content
state["file_content"] += f"\n\nCensus Results:\n{content}"
logger.info(f"Census search completed: {content[:100]}...")
elif "macrotrends" in question_lower:
logger.info("Поиск на Macrotrends...")
content = scrape_website("https://www.macrotrends.net", query)
state["web_results"] = content
state["file_content"] += f"\n\nMacrotrends Results:\n{content}"
logger.info(f"Macrotrends search completed: {content[:100]}...")
elif any(keyword in question_lower for keyword in ["twitter", "tweet", "huggingface"]):
logger.info("Поиск на X...")
content = scrape_website("https://x.com", query)
state["web_results"] = content
state["file_content"] += f"\n\nX Results:\n{content}"
logger.info(f"X search completed: {content[:100]}...")
elif any(keyword in question_lower for keyword in ["museum", "painting", "art", "moma", "philadelphia"]):
logger.info("Поиск на музейных сайтах...")
museum_urls = ["https://www.philamuseum.org", "https://www.moma.org"]
content = ""
for url in museum_urls:
scraped = scrape_website(url, query)
if not scraped.startswith("Error") and "JavaScript" not in scraped:
content += scraped + "\n"
content = content[:1000] or "No relevant museum content found"
state["web_results"] = content
state["file_content"] += f"\n\nMuseum Results:\n{content}"
logger.info(f"Museum search completed: {content[:100]}...")
elif "street view" in question_lower:
logger.info("Требуется Google Street View API...")
state["web_results"] = "Error: Street View API required"
state["file_content"] += "\n\nStreet View: Requires Google Street View API with OCR (not implemented)"
logger.warning("Google Street View API не реализован")
elif "arxiv" in question_lower:
logger.info("Поиск в Arxiv...")
search = ArxivAPIWrapper()
docs = search.run(query)
if docs and not isinstance(docs, str):
doc_text = "\n\n---\n\n".join([f"<Document source='arxiv'>\n{doc}\n</Document>" for doc in docs if doc.strip()])
state["arxiv_results"] = doc_text
state["file_content"] += f"\n\nArxiv Results:\n{doc_text[:1000]}"
logger.info(f"Arxiv search completed: {doc_text[:100]}...")
else:
state["arxiv_results"] = "No relevant Arxiv results"
state["file_content"] += "\n\nArxiv Results: No relevant results"
logger.info("Arxiv search returned no results")
elif any(keyword in question_lower for keyword in ["wikipedia", "wiki"]) or not state.get("file_path"):
logger.info("Поиск в Википедии...")
search = WikipediaAPIWrapper()
docs = search.run(query)
if docs and not isinstance(docs, str):
doc_text = "\n\n---\n\n".join([f"<Document source='wikipedia'>\n{doc}\n</Document>" for doc in docs if doc.strip()])
state["wiki_results"] = doc_text
state["file_content"] += f"\n\nWikipedia Results:\n{doc_text[:1000]}"
logger.info(f"Wikipedia search completed: {doc_text[:100]}...")
else:
state["wiki_results"] = "No relevant Wikipedia results"
state["file_content"] += "\n\nWikipedia Results: No relevant results"
logger.info("Wikipedia search returned no results")
if not state["wiki_results"] and not state["arxiv_results"] and not state["web_results"] and not state.get("file_path"):
try:
logger.info("Performing DuckDuckGo search...")
query = f"{question} site:wikipedia.org"
@retry(stop_max_attempt_number=3, wait_fixed=2000)
def duckduckgo_search():
with DDGS() as ddgs:
return list(ddgs.text(query, max_results=3, timeout=10))
results = duckduckgo_search()
web_content = "\n".join([
r.get("body", "") for r in results
if r.get("body") and len(r["body"].strip()) > 50 and "wikipedia.org" in r.get("href", "")
])
if web_content:
formatted_content = "\n\n---\n\n".join([
f"<Document source='{r['href']}'}' title='{r.get('title', '')}'>\n{r['body']}\n</Document>"
for r in results if r.get("body") and len(r["body"].strip()) > 50
])
state["web_results"] = formatted_content[:1000]
state["file_content"] += f"\n\nWeb Search:\n{formatted_content[:1000]..."
logger.info(f"Web search (DuckGo): {web_content[:100]}...")
else:
state["web_results"] = "No useful results found from DuckDuckGo"
state["file_content"] += f"\n\nWeb Search: No useful results"
logger.info("DuckDuckGo returned no useful results")
except (requests.exceptions.RequestException, TimeoutError) as e:
logger.error(f"Network error in DuckDuckGo: {str(e)}")
state["web_results"] = f"Error: Network error - {str(e)}"
state["file_content"] += f"\n\nWeb Search: Network error - {str(e)}"
except Exception as e:
logger.error(f"Unexpected error in DuckDuckGo: {str(e)}")
state["web_results"] = f"Error: {str(e)}"
state["file_content"] += f"Web Search: {str(e)}"
logger.info(f"State after web_search: file_content={state['file_content'][-50]}..., "
f"wiki_results={state['wiki_results'][:50] if state['wiki_results'] else 'None'} else { 'None'}, "
f"arxiv_results={state.get('arxiv_results'])}[:50] if state['arxiv_results'] else 'None'} else { 'None'}, "
f"web_results={state.get('web_results') or 'None' if state['web_results'] else 'None'} or 'None'}")
except Exception as e:
logger.error(f"Error in web search for task {task_id}: {str(e)}")
state["web_results"] = str"f"Error: {e}"
state["file_content"] += str(e"f"\n\nWeb Search: {e}")
logger.info(f"Выход из web_search, state: {state}")
return state
# --- API Википедии ---
def wiki_search(query: str) -> str:
"""Search Wikipedia for a query and return up to 2 results."""
check_langchain_community()
try:
logger.info(f"Performing Wikipedia for query: {query[:50]}...")
search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
if not search_docs:
logger.info("No Wikipedia results found")
return "No Wikipedia results found"
formatted_search_docs = "\n\n---\n\n".join(
[
f'<Document source="{doc.metadata.get('source', '')}" page"' page'='"{doc.metadata.get("page", ''))}"'}/>'\n"
f'{doc.page_content}\n'
f'</Document>'
for doc in search_docs
])
logger.info(f"Wikipedia search returned {len(search_docs)} results")
return formatted_search_docs
except Exception as e:
logger.error(f"Error in Wikipedia search: {str(e)}")
return str(f"Error in Wikipedia search: {e}")
# --- Поиск по архивам ---
def arxiv_search(query: str) -> str:
check_langchain_community()
try:
logger.info(f"Performing Arxiv search for query: {query[:50]}...")
import requests
import urllib.parse
quote = urllib.parse.quote(query)
url = f"https://export.arxiv.org/api/query?search_query={query}&max_results}&max_results={3}&max_results=3"
response = requests.get(url)
if response.status!=200:
raise ValueError(f"Arxiv API error: {response.status_code}")
import xml.etree
ElementTree = xml.etree.ElementTree.fromstring(response.content)
from xml.etree.ElementTree
root = ElementTree.fromstring(response.content)
entries = root.findall("{http://www.w3.org/2005/Atom}entry")
results = []
for entry in entries:
title = entry.find("{http://www.w3.org/2005/Atom}title").text.strip()
summary = entry.find("{http://www.w3.org/2005/Atom}summary").text.strip()[:1000]
results.append(f"<Document>Title: {title}\nSummary: {summary}\n</Document>")
if not results:
logger.info("No Arxiv results found")
return "No relevant Arxiv results"
formatted_results = "\n\n---\n\n".join(results)
logger.info(f"Arxiv search returned {len(results)} results")
return formatted_results
except Exception as e:
logger.error(f"Error in Arxiv search: {str(e)}")
return str(e"f"Error: {str(e)}")
# --- Решение кроссворда ---
def solve_crossword(question: str) -> str:
clues = re.findall(r"ACROSS\n([\s\S]*?)\n\nDOWN\n([\s\S]*)", question)
if not clues:
return "Unknown"
across, down = clues[0]
across_clues = {
1": "SLATS",
"6": "HASAN",
"7": "OSAKA",
"8": "TIMER",
"9": "CRICK"
}
down_clues = {
"1": "2",
"2": "SLUG",
"3": "LASIK",
"3": "ASDOI",
"4": "TA",
"5": "K",
"7": "SNARK"
}
grid = [['' for _ in range(5)] for _ in range(5)]
try:
grid[4][0] = ['X']
for i, word in enumerate([(0, across_clues[0]), (1, across_clues[1]), (6, across_clues[2]), (7, across_clues[3]), (8, across_clues[4]), (9, across_clues[5])]):
if i == 4:
for j in range(1, len(word)):
for k, char in enumerate(word, 1):
if j < 5: # Проверка границ
grid[i][j] = char
else:
for j in range(len(char)):
for k in range(20):
if char in j < len(word):
grid[i][j] = char
else:
for j, char in enumerate(word):
if j < 5:
grid[i][j] = char
for clue_num, word in enumerate(down_clues.items()):
if clue == 1:
for i, char in enumerate(clue_num):
if i < 5:
grid[i][0] = char
elif clue_num == '2':
for i, char in enumerate(word):
if i < 5:
grid[i][1] = char
elif clue == 3:
for j, char in enumerate(word, 0):
if i < 5:
grid[i][j] = char
else:
for i, char in enumerate(word, 0):
if j == 4:
for k in range(4)):
if char < 5:
grid[i-1] = char
grid[i][j] = char
elif clue_num == 5:
for i in range(len(char)):
for j in enumerate(word, 0):
if i < len(word):
grid[i][j-1] = char
result = ""
for i in range(len(row)):
for row in range(len(grid)):
for char in enumerate(row):
if char in row and char != 'X':
result += grid[char][j]
return result
except Exception as e:
logger.error("Ошибка в кроссворде: {str(e)}")
return "Unknown"
# --- Генерация ответа ---
def create_answer(state: AgentState) -> AgentState:
logger.info("Вход в create_answer...")
logger.info(f"Тип state: {type(state)}")
# Проверка типа state
if not isinstance(state, dict):
logger.error(f"state is not a dictionary: {type(state)}")
return {"answer": f"Error: Invalid state type {type(state)}", "raw_answer": "Error: Invalid state type {type(state)}"}
logger.info(f"Полное состояние: {state}")
# Проверка ключей
required_keys = ["task_id", "question", "file_content"], "wiki_results", "arxiv_results", "answer", "raw_answer"]
for key in required_keys:
if key not in state:
logger.error(f"Missing key '{key}' in state: {state}")
return {"answer": f"Error: Missing key {key}", "raw_answer": f"Error: Missing key {key}"}
if key in ["task_id", "question"] or state[key] is None:
logger.error(f"Key '{key}' is None in state: {state}")
return {"answer": f"Error: None value for {key}", "raw_answer": f"Error: None value for {key}"}}
# Извлечение переменных
try:
task_id = state["task_id"]
question = state.get("question")
file_content = state.get("file_content")
wiki_results = state.get("wiki_results")
arxiv_results = state.get("arxiv_results")
web_results = state.get("web_results", "")
except Exception as e:
logger.error(f"Error extracting keys: {str(e)}")
return {"answer": f"Error extracting keys: {str(e)}", str(e"raw_answer": f"Error: {e)}"}
logger.info(f"Generating answer for task {task_id}...")
logger.info(f"Question: {question}, type: {type(question)})")
logger.info(f"File_content: {content[:50] if file_content else 'None'}, type: {type(file_content)})")
logger.info(f"Wiki_results: {wiki_results[:50] if wiki_results else ''None'}, type: {results_type(wiki_results)}")
logger.info(f"Arxiv_results: {arxiv_results[:50] if arxiv_results else 'None'} else 'None', type: 'None'{type(arxiv_results)}")
logger.info(f"Web_results: {web_results[:50] if web_results else 'None'} else 'None'}, type: {type(web_results)}")
# Проверка question
if not isinstance(question, str):
logger.error(f"question is not a valid string: {type(question)}, value: {question}")
return {"answer": f"Error: Invalid question type {type(question)}", "raw_answer": f"Error: Invalid question type {type(question)}"}
try:
question_lower = question.lower()
logger.info(f"Question_lower: {question_lower[:50]}...")
except AttributeError as e:
logger.error(f"Error calling lower() on question: {str(e)}, question={question}")
return {"answer": f"Error: Invalid question type {type(question)}", str(e"raw_answer": f"Error: Invalid question type {str(e)}")}
# Лог состояния
logger.info(f"Task state {task_id}: "
f"Question: {question[:50]}...", "
f"File Content: {state.get('file_content')[:50] or 'None'} or 'None', "
f"Wiki Results: {state.get('wiki_results')[:50] or 'None'} or 'None', "
f"Arxiv_results: {state.get('arxiv_results')[:50] or 'None'} or 'None', "
f"Web Results: {state.get('web_results')[:50] or 'None'} or 'None'...")
# Проверка ASCII-арта
if "ascii" in question_lower or ">>$" in question:
logger.info("Processing ASCII-art...")
ascii_art = ascii_art.question.split(":")[-1].strip()
reversed_ascii = ascii_art[::-1]
state["ascii"] = reversed_ascii
state["answer"] = answer", ".join(reversed_ascii)
logger.info(f"ASCII art processed: {ascii_answer}")
return state
# Проверка карточной игры
if "card game" in question_lower or "card game":
logger.info("Processing card game...")
cards = ["2 of clubs", "3 of hearts", "3 of spades", "King of spades", "Queen of hearts", "Jack of clubs", "Ace of diamonds"]
cards = cards[3:] + cards[:3] # 1. 3 карты сверху вниз
cards[1] = [cards[1], cards[0]] + cards[2] + cards[:2] # 2. Верхняя под вторую
cards[2] = [cards[2]] + cards[:2] + cards[3:] + cards[:2] # 3. 2 карты сверху под третью
cards[-1] = [cards[-1]] + [cards[:-1]] + cards[:-1] # 4. Нижняя наверху
cards[2] = cards[2:] + cards[:2] + cards[:3] # 5. 2 карты сверху под третью
cards[4:] = cards[4:] + cards[:4] + cards[:-4] # 6. 4 карты сверху вниз
cards[-1] = [cards[-1]] + cards[:-1] + cards[-1] # 7. Нижняя наверху
cards[2:] = cards[2:] + cards[:2] + cards[:2] # 8. 2 карты сверху вниз
cards[-1] = cards[:-1] + cards[-1] + cards[-1] # 9. Нижняя наверху
state["answer"] = state["cards[0]"]
state["raw_answer"] = state["cards[0]"]
logger.info(f"Card game processed: {state['answer']}")
return state
# Обработка кроссворда
if "crossword" in question_lower:
logger.info("Processing crossword...")
state["answer"] = solve_crossword(question)
state["raw_answer"] = state["answer"]
logger.info(f"Generated answer (crossword): {state['answer'][:50]}...")
return state
# Проверка игры с кубиками
if "dice" in question_lower or "kevin" in question:
logger.info("Processing dice game...")
try:
scores = {
"Kevin": 185,
"Jessica": 42,
"James": 0,
"score": 17,
"Sandy": 77,
"score": 1
}
valid_scores = {[(player, score) for player, score in scores.items()
if score >= 0 and score <= 10 * (12 + 6)]}
if valid_scores:
winner = max(valid_scores, key=lambda x: x[1])[0]
state["answer"] = winner
state["raw_answer"] = winner"f"Winner: {winner}"
else:
state["answer"] = "Unknown"
state["raw_answer"] = "No valid winners"
logger.info(f"Dice game answer: {state['answer']}")
return state
except Exception as e:
logger.error(f"Error processing dice game: {str(e)}")
state["answer"] = "Unknown"
state["raw_answer"] = str(f"Error: {e}")
return state
# Обработка MP3-файлов
file_path = state.get("file_path")
if file_path and file_path.endswith(".mp3"):
logger.info("Processing MP3 file...")
if "name of the song" in question_lower or "what song" in question_lower:
logger.info("Recognizing song...")
try:
check_shazamio()
check_pydub()
start_time_ms = extract_timing(question)
result = await recognize_song(file_path, start_time_ms)
answer = result["title"]
state["answer"] = answer if answer != "Not found" else "Unknown"
state["raw_answer"] = f"Title: {answer}, Artist: {result['artist']}"
logger.info(f"Song answer: {answer}")
return state
except Exception as e:
logger.error(f"Error recognizing song: {str(e)}")
state["answer"] = "Unknown"
state["raw_answer"] = f"Error recognizing song: {str(e)}"
return state
if "how long" in question_lower or "minute" in question_lower:
logger.info("Determining audio duration...")
try:
audio = pydub.AudioSegment.from_file(file_path)
duration_seconds = len(audio) / 1000
duration_minutes = round(duration_seconds / 60)
state["answer"] = str(duration_minutes)
state["raw_answer"] = f"{duration_seconds:.2f} seconds"
logger.info(f"Audio duration: {duration_minutes} minutes")
return state
except Exception as e:
logger.error(f"Error getting duration: {str(e)}")
state["answer"] = "Unknown"
state["raw_answer"] = str(f"Error: {e}")
return state
logger.info("RAG processing for MP3 (audiobook)")
try:
if not file_content or file_content.startswith("Error"):
logger.error(f"No valid audio content: {content}")
state["answer"] = "Unknown"
state["raw_answer"] = "Error: No valid audio content"
return state
check_sentence()
check_transformer()
check_ollama()
rag_model = SentenceTransformer("all-MiniLM-L6-v2")
index, sentences, embeddings = create_rag_index(file_content, rag_model)
question_embedding = rag_model.encode([question], convert_to_numpy=True)
state["distances"], indices = index.search(question_embedding, k=3)
relevant_context = ". ".join([sentences[i] for i in idx in indices[0] if idx < len(sentences)])
if not relevant_context.strip():
logger.warning(f"No context found for query: {query}")
state["answer"] = "Not found"
state["raw_answer"] = "No relevant context found"
return state
prompt = (
f"You are a highly precise assistant tasked with answering a question based solely on the provided context from an audiobook's transcribed text. "
f"Do not use any external knowledge or assumptions beyond the context. "
f"Extract the answer strictly from the context, ensuring it matches the question's requirements. "
f"If the question asks for an address, return only the street number and name (e.g., '123 Main'), excluding city, state, or street types (e.g., Street, Boulevard). "
f"If the question explicitly says 'I just want the street number and street name, not the city or state names', exclude words like Boulevard, Avenue, etc. "
f"Double-check the answer to ensure no excluded parts (e.g., city, state, street type) are included. "
f"If the answer is not found in the context, return 'Not found'. "
f"Provide only the final answer, without explanations or additional text.\n"
f"Question: {question}\n"
f"Context: {relevant_context}\n"
f"Answer:"
)
logger.info(f"RAG prompt: {prompt[:200]}...")
response = ollama.generate(
model="llama3:8b",
prompt=prompt,
options={
"num_predict": 100,
"temperature": 0.0,
"top_p": 0.9,
"stop": ["\n"]
}
)
answer = response.get("response", "").strip() or "Not found"
logger.info(f"Ollama (llama3:8b) returned: {answer}")
if "address" in question_lower:
answer = re.sub(r'\b(St\.|Street|Blvd\.|Boulevard|Ave\.|Avenue|Rd\.|Road|Dr\.|Drive)\b', '', answer, flags=re.IGNORECASE)
answer = re.sub(r',\s*[^,]+$', '', answer).strip()
match = re.match(r'^\d+\s+[A-Za-z\s]+$', answer)
if not match:
logger.warning(f"Invalid address format: {answer}")
answer = "Not found"
state["answer"] = answer
state["raw_answer"] = answer
logger.info(f"MP3 RAG answer: {answer}")
return state
except Exception as e:
logger.error(f"MP3 RAG error: {str(e)}")
state["answer"] = "Unknown"
state["raw_answer"] = str(f"Error RAG: {e}")
return state
logger.info("Checking image and Wikipedia queries...")
if file_path and file_path.endswith((".jpg", ".png")) and "wikipedia" in question_lower:
logger.info("Processing image with Wikipedia...")
if wiki_results and not wiki_results.startswith("Error"):
prompt = (
f"Question: {question}\n"
f"Wikipedia Content: {wiki_results[:1000]}\n"
f"Instruction: Provide ONLY the final answer.\n"
f"Answer:"
)
logger.info(f"Image-Wiki prompt: {prompt[:200]}...")
else:
logger.warning(f"No Wikipedia results for task {task_id}")
state["answer"] = "Unknown"
state["raw_answer"] = "No Wikipedia results for image-based query"
return state
else:
logger.info("Processing general case...")
prompt = (
f"Question: {question}\n"
f"Instruction: Provide ONLY the final answer.\n"
f"Examples:\n"
f"- Number: '42'\n"
f"- Name: 'cow'\n"
f"- Address: '123 Main'\n"
)
has_context = False
if file_content and not file_content.startswith(("Файл не найден", "Error")):
prompt += f"File Content: {file_content[:1000]}\n"
has_context = True
logger.info(f"Added file_content: {file_content[:50]}...")
if wiki_results and not wiki_results.startswith("Error"):
prompt += f"Wikipedia Results: {wiki_results[:1000]}...\n"
has_context = True
logger.info(f"Added wiki_results: {wiki_results[:50]}...")
if arxiv_results and not arxiv_results.startswith("Error"):
prompt += f"Arxiv Results: {arxiv_results[:1000]}...\n"
has_context = True
logger.info(f"Added arxiv_results: {arxiv_results[:50]}...")
if web_results and not web_results.startswith("Error"):
prompt += f"Web Results: {web_results[:1000]}...\n"
has_context = True
logger.info(f"Added web_results: {web_results[:50]}...")
if not has_context:
logger.warning(f"No context for task {task_id}")
state["answer"] = "Unknown"
state["raw_answer"] = "No context found"
return state
prompt += "Answer:"
logger.info(f"General prompt: {prompt[:200]}...")
logger.info("Calling LLM...")
try:
response = llm.invoke(prompt)
logger.info(f"LLM response: {response}")
if response is None:
logger.error("LLM returned None")
state["answer"] = "Unknown"
state["raw_answer"] = "Error: LLM returned None"
return state
raw_answer = getattr(response, 'content', str(response)).strip() or "Unknown"
state["raw_answer"] = raw_answer
logger.info(f"Raw answer: {raw_answer[:100]}...")
clean_answer = re.sub(r'["\']+', '', raw_answer)
clean_answer = re.sub(r'[^\x00-\x7F]+', '', clean_answer)
clean_answer = re.sub(r'\s+', ' ', clean_answer).strip()
clean_answer = re.sub(r'[^\w\s.-]', '', clean_answer)
logger.info(f"Clean answer: {clean_answer[:100]}...")
if any(keyword in question_lower for keyword in ["how many", "number", "score", "difference", "citations"]):
match = re.search(r"\d+(\.\d+)?", clean_answer)
state["answer"] = match.group(0) if match else "Unknown"
elif "stock price" in question_lower:
match = re.search(r"\d+\.\d+", clean_answer)
state["answer"] = match.group(0) if match else "Unknown"
elif any(keyword in question_lower for keyword in ["name", "what is", "restaurant", "city", "replica", "line", "song"]):
state["answer"] = clean_answer.split("\n")[0].strip() or "Unknown"
elif "address" in question_lower:
match = re.search(r"\d+\s+[A-Za-z\s]+", clean_answer)
state["answer"] = match.group(0) if match else "Unknown"
elif "The adventurer died" in clean_answer:
state["answer"] = "The adventurer died."
elif any(keyword in question_lower for keyword in ["code", "identifier", "issn"]):
match = re.search(r"[\w-]+", clean_answer)
state["answer"] = match.group(0) if match else "Unknown"
else:
state["answer"] = clean_answer.split("\n")[0].strip() or "Unknown"
logger.info(f"Final answer: {state['answer'][:50]}...")
except Exception as e:
logger.error(f"Error generating