First Version
This commit is contained in:
aliakbar
@@ -0,0 +1,33 @@
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install system dependencies
|
||||
RUN apt-get update && apt-get install -y \
|
||||
gcc \
|
||||
g++ \
|
||||
build-essential \
|
||||
curl \
|
||||
wget \
|
||||
git \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Copy requirements first for better caching
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir --upgrade pip && \
|
||||
pip install --no-cache-dir wheel setuptools && \
|
||||
pip install --no-cache-dir -r requirements.txt && \
|
||||
pip install fastapi==0.111.0
|
||||
pip install torch==2.1.0 --index-url https://download.pytorch.org/whl/cpu
|
||||
|
||||
# Copy application code
|
||||
COPY . .
|
||||
|
||||
# Create necessary directories
|
||||
RUN mkdir -p /app/faiss_index /root/.cache/huggingface
|
||||
|
||||
# Expose port
|
||||
EXPOSE 8000
|
||||
|
||||
# Run the application
|
||||
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--reload"]
|
||||
@@ -0,0 +1,371 @@
|
||||
from typing_extensions import TypedDict, List
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from pydantic import BaseModel, Field
|
||||
from langgraph.store.memory import InMemoryStore
|
||||
from langmem import create_memory_store_manager
|
||||
from langchain_openai import ChatOpenAI
|
||||
import json
|
||||
from pathlib import Path
|
||||
import asyncio
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from functools import lru_cache
|
||||
import pickle
|
||||
from dataclasses import dataclass
|
||||
from queue import Queue
|
||||
import logging
|
||||
|
||||
from rag.rag import RAG
|
||||
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class Episode(BaseModel):
|
||||
observation: str = Field(..., description="The context and setup - what happened")
|
||||
thoughts: str = Field(
|
||||
...,
|
||||
description="Internal reasoning process and observations of the agent in the episode that let it arrive"
|
||||
' at the correct action and result. "I ..."',
|
||||
)
|
||||
action: str = Field(
|
||||
...,
|
||||
description="What was done, how, and in what format. (Include whatever is salient to the success of the action). I ..",
|
||||
)
|
||||
result: str = Field(
|
||||
...,
|
||||
description="Outcome and retrospective. What did you do well? What could you do better next time? I ...",
|
||||
)
|
||||
|
||||
class StateAgent(TypedDict):
|
||||
query: List[List[str]]
|
||||
top_k: int
|
||||
need_correction: bool
|
||||
is_correct: str
|
||||
|
||||
class AsyncBackgroundProcessor:
|
||||
"""Background processor for async operations"""
|
||||
def __init__(self, max_workers=2):
|
||||
self.executor = ThreadPoolExecutor(max_workers=max_workers)
|
||||
self.store_queue = Queue()
|
||||
self.running = True
|
||||
self.worker_thread = threading.Thread(target=self._process_queue, daemon=True)
|
||||
self.worker_thread.start()
|
||||
|
||||
def _process_queue(self):
|
||||
"""Process store operations in background"""
|
||||
while self.running:
|
||||
try:
|
||||
if not self.store_queue.empty():
|
||||
func, args, kwargs = self.store_queue.get(timeout=0.1)
|
||||
try:
|
||||
func(*args, **kwargs)
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing background task: {e}")
|
||||
else:
|
||||
time.sleep(0.01)
|
||||
except Exception as e:
|
||||
logger.error(f"Queue processing error: {e}")
|
||||
|
||||
def submit(self, func, *args, **kwargs):
|
||||
"""Submit task to background queue"""
|
||||
self.store_queue.put((func, args, kwargs))
|
||||
|
||||
def shutdown(self):
|
||||
"""Shutdown background processor"""
|
||||
self.running = False
|
||||
self.executor.shutdown(wait=True)
|
||||
if self.worker_thread.is_alive():
|
||||
self.worker_thread.join(timeout=5)
|
||||
|
||||
class Agent:
|
||||
def __init__(self, Config):
|
||||
if not Config.API_KEY:
|
||||
raise RuntimeError("Missing API_KEY. Set it in your .env or environment as OPENAI_API_KEY/API_KEY.")
|
||||
os.environ.setdefault("OPENAI_API_KEY", Config.API_KEY)
|
||||
os.environ.setdefault("OPENAI_BASE_URL", "https://api.avalai.ir/v1")
|
||||
|
||||
# Initialize components
|
||||
self.rag = RAG(
|
||||
Config.MODEL_NAME,
|
||||
Config.RETREIVER_NAME,
|
||||
Config.VECTOR_PATH,
|
||||
Config.DATA_PATH,
|
||||
Config.DEVICE
|
||||
)
|
||||
self.memory_llm = self.rag.generator
|
||||
|
||||
# Load store from file first to avoid duplicate loading
|
||||
self.dump_file = Config.DUMP_FILE if hasattr(Config, 'DUMP_FILE') else "episodes_dump.json"
|
||||
self.episode_cache = {} # Cache for episodes
|
||||
self.search_cache = {} # Cache for search results
|
||||
|
||||
# Initialize store with optimized settings
|
||||
self.store = InMemoryStore(
|
||||
index={
|
||||
"dims": 1536,
|
||||
"embed": "openai:text-embedding-3-small",
|
||||
"metric": "cosine", # Use cosine similarity
|
||||
}
|
||||
)
|
||||
|
||||
# Background processor for async operations
|
||||
self.bg_processor = AsyncBackgroundProcessor(max_workers=2)
|
||||
|
||||
# Try to load store from file
|
||||
if Path(self.dump_file).exists():
|
||||
try:
|
||||
self._fast_load_store(filepath=self.dump_file)
|
||||
logger.info(f"Loaded {len(self.episode_cache)} episodes from cache")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not load previous content: {e}")
|
||||
|
||||
# Initialize chat with optimized settings
|
||||
chat = ChatOpenAI(
|
||||
model=self.memory_llm.model_name,
|
||||
openai_api_key=os.environ["OPENAI_API_KEY"],
|
||||
base_url=os.environ.get("OPENAI_BASE_URL", "https://api.avalai.ir/v1"),
|
||||
temperature=0.2,
|
||||
max_tokens=512,
|
||||
timeout=30, # Add timeout
|
||||
max_retries=2, # Limit retries
|
||||
)
|
||||
|
||||
# Initialize memory manager
|
||||
self.manager = create_memory_store_manager(
|
||||
chat,
|
||||
namespace=("memories", "episodes"),
|
||||
schemas=[Episode],
|
||||
instructions="Extract exceptional examples of noteworthy problem-solving scenarios",
|
||||
enable_inserts=True,
|
||||
store=self.store,
|
||||
)
|
||||
|
||||
# Track store size for periodic dumping
|
||||
self.operations_since_last_dump = 0
|
||||
self.dump_threshold = 10 # Dump every 10 operations
|
||||
|
||||
def _fast_load_store(self, filepath: str = "episodes_dump.json"):
|
||||
"""Fast loading of store data with caching"""
|
||||
try:
|
||||
with open(filepath, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
# Bulk load episodes into cache
|
||||
for key, episode_data in data.items():
|
||||
if 'content' in episode_data['value']:
|
||||
self.episode_cache[key] = episode_data['value']['content']
|
||||
|
||||
logger.info(f"Loaded {len(self.episode_cache)} episodes into cache")
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading store: {e}")
|
||||
raise
|
||||
|
||||
@lru_cache(maxsize=100)
|
||||
def _cached_search(self, query: str, limit: int):
|
||||
"""Cached search function"""
|
||||
cache_key = f"{query[:50]}_{limit}"
|
||||
|
||||
if cache_key in self.search_cache:
|
||||
return self.search_cache[cache_key]
|
||||
|
||||
# Perform actual search
|
||||
results = self.store.search(
|
||||
("memories", "episodes"),
|
||||
query=query,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
# Cache the results
|
||||
self.search_cache[cache_key] = results
|
||||
return results
|
||||
|
||||
def _generate_system_message(self, similar_episodes):
|
||||
"""Generate system message from similar episodes"""
|
||||
system_message = "You are a helpful assistant."
|
||||
|
||||
if similar_episodes:
|
||||
system_message += "\n\n### EPISODIC MEMORY:"
|
||||
for i, item in enumerate(similar_episodes, start=1):
|
||||
episode = item.value["content"]
|
||||
system_message += f"\n\nEpisode {i}:"
|
||||
system_message += f"\nObservation: {episode.get('observation', 'N/A')}"
|
||||
system_message += f"\nThought: {episode.get('thoughts', 'N/A')}"
|
||||
system_message += f"\nAction: {episode.get('action', 'N/A')}"
|
||||
system_message += f"\nResult: {episode.get('result', 'N/A')}"
|
||||
|
||||
return system_message
|
||||
|
||||
def app(self, messages: list, top_k: int, feedback: bool = False):
|
||||
"""Optimized app method with async processing"""
|
||||
logger.info('Starting processing...')
|
||||
|
||||
# Get user message
|
||||
user_message = messages[-1]["content"]
|
||||
|
||||
# 1. Search for similar episodes
|
||||
logger.info('Searching for similar episodes...')
|
||||
similar = self._cached_search(user_message, top_k)
|
||||
|
||||
# 2. Generate response
|
||||
logger.info('Generating response...')
|
||||
system_message = self._generate_system_message(similar)
|
||||
|
||||
# Use thread pool for async generation
|
||||
def generate_response():
|
||||
return self.memory_llm.generate([
|
||||
{"role": "system", "content": system_message},
|
||||
*messages[:-1],
|
||||
{"role": "user", "content": user_message}
|
||||
])
|
||||
|
||||
# Run in thread pool to avoid blocking
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(generate_response)
|
||||
response = future.result(timeout=30) # 30 second timeout
|
||||
|
||||
logger.info('Response generated')
|
||||
|
||||
# 3. Prepare messages for background processing
|
||||
messages_with_assistant = [
|
||||
*messages,
|
||||
{"role": "assistant", "content": response}
|
||||
]
|
||||
|
||||
# 4. Submit store operation to background processor (non-blocking)
|
||||
if feedback:
|
||||
self.bg_processor.submit(
|
||||
self._process_store_operation,
|
||||
messages_with_assistant,
|
||||
user_message,
|
||||
response
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
def _process_store_operation(self, messages_with_assistant, user_message, response):
|
||||
"""Process store operation in background"""
|
||||
try:
|
||||
logger.info('Processing store operation in background...')
|
||||
self.manager.invoke({"messages": messages_with_assistant})
|
||||
logger.info('Store operation completed')
|
||||
|
||||
# Increment operation counter
|
||||
self.operations_since_last_dump += 1
|
||||
|
||||
# Periodic dumping
|
||||
if self.operations_since_last_dump >= self.dump_threshold:
|
||||
self._background_dump_store()
|
||||
self.operations_since_last_dump = 0
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Store operation failed: {e}")
|
||||
# Fallback: add correction directly
|
||||
self._add_correction_direct(
|
||||
user_message,
|
||||
response
|
||||
)
|
||||
|
||||
def _add_correction_direct(self, observation: str, result: str):
|
||||
"""Add correction directly without manager"""
|
||||
episode = Episode(
|
||||
observation=observation,
|
||||
thoughts="Saved via fallback due to upstream error.",
|
||||
action="Generated assistant reply and stored episode directly.",
|
||||
result=result,
|
||||
)
|
||||
key = f"episode:{uuid.uuid4()}"
|
||||
self.store.put(
|
||||
("memories", "episodes"),
|
||||
key=key,
|
||||
value={"content": episode.model_dump()}
|
||||
)
|
||||
|
||||
# Update cache
|
||||
self.episode_cache[key] = episode.model_dump()
|
||||
|
||||
# Clear search cache since we added new data
|
||||
self.search_cache.clear()
|
||||
|
||||
# Schedule background dump
|
||||
self.bg_processor.submit(self._background_dump_store)
|
||||
|
||||
def _background_dump_store(self, filepath: str = None):
|
||||
"""Dump store to file in background"""
|
||||
if filepath is None:
|
||||
filepath = self.dump_file
|
||||
|
||||
try:
|
||||
# Get all items efficiently
|
||||
items = []
|
||||
batch_size = 50
|
||||
|
||||
# Try to get items in batches if store supports it
|
||||
try:
|
||||
# Try batch retrieval if available
|
||||
items = list(self.store.search(
|
||||
("memories", "episodes"),
|
||||
query="", # Empty query to get all
|
||||
limit=1000 # Limit to prevent memory issues
|
||||
))
|
||||
except:
|
||||
# Fallback to cache
|
||||
items = [{
|
||||
'key': key,
|
||||
'namespace': ("memories", "episodes"),
|
||||
'value': {'content': value}
|
||||
} for key, value in self.episode_cache.items()]
|
||||
|
||||
# Prepare data for JSON
|
||||
all_data = {}
|
||||
for item in items:
|
||||
if hasattr(item, 'key'):
|
||||
key = item.key
|
||||
value = item.value
|
||||
else:
|
||||
key = item['key']
|
||||
value = item['value']
|
||||
|
||||
all_data[key] = {
|
||||
"namespace": ("memories", "episodes"),
|
||||
"value": value
|
||||
}
|
||||
|
||||
# Write to file
|
||||
temp_file = f"{filepath}.tmp"
|
||||
with open(temp_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(all_data, f, indent=2, ensure_ascii=False)
|
||||
|
||||
# Atomic replace
|
||||
os.replace(temp_file, filepath)
|
||||
|
||||
logger.info(f"Successfully dumped {len(all_data)} episodes")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error dumping store: {e}")
|
||||
|
||||
def search_episodes(self, query: str, limit: int = 5):
|
||||
"""Optimized search method"""
|
||||
return self._cached_search(query, limit)
|
||||
|
||||
def shutdown(self):
|
||||
"""Clean shutdown"""
|
||||
# Dump store before shutdown
|
||||
self._background_dump_store()
|
||||
|
||||
# Shutdown background processor
|
||||
self.bg_processor.shutdown()
|
||||
|
||||
logger.info("Agent shutdown complete")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from utils.config import Config
|
||||
|
||||
agent = Agent(Config)
|
||||
result = agent.app([{"role": "user", "content": "hi how are you?"}], 5)
|
||||
print(f"Response: {result}")
|
||||
|
||||
agent.shutdown()
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,16 @@
|
||||
from utils.app import Application
|
||||
import signal
|
||||
import threading
|
||||
|
||||
|
||||
def main():
|
||||
app = Application()
|
||||
try:
|
||||
app.run_fastapi()
|
||||
except KeyboardInterrupt:
|
||||
print("\nFastAPI server interrupted.")
|
||||
finally:
|
||||
print("Application shutdown complete.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,148 @@
|
||||
from datetime import datetime
|
||||
from flask_sqlalchemy import SQLAlchemy
|
||||
from flask_login import UserMixin
|
||||
from flask_bcrypt import Bcrypt
|
||||
|
||||
db = SQLAlchemy()
|
||||
bcrypt = Bcrypt()
|
||||
|
||||
class User(UserMixin, db.Model):
|
||||
__tablename__ = 'users'
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
username = db.Column(db.String(80), unique=True, nullable=False)
|
||||
email = db.Column(db.String(120), unique=True, nullable=False)
|
||||
password_hash = db.Column(db.String(255), nullable=False)
|
||||
|
||||
# فیلدهای جدید برای پنل ادمین
|
||||
role = db.Column(db.String(20), default='user') # 'admin', 'user'
|
||||
is_active = db.Column(db.Boolean, default=True)
|
||||
created_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
|
||||
def set_password(self, password):
|
||||
"""Hash and set the password"""
|
||||
self.password_hash = bcrypt.generate_password_hash(password).decode('utf-8')
|
||||
|
||||
def check_password(self, password):
|
||||
"""Check if the provided password matches the hash"""
|
||||
return bcrypt.check_password_hash(self.password_hash, password)
|
||||
|
||||
def __repr__(self):
|
||||
return f'<User {self.username}>'
|
||||
|
||||
class DataFile(db.Model):
|
||||
__tablename__ = 'data_files'
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
filename = db.Column(db.String(255), nullable=False)
|
||||
original_filename = db.Column(db.String(255), nullable=False)
|
||||
file_path = db.Column(db.String(500), nullable=False)
|
||||
file_size = db.Column(db.Integer, nullable=False)
|
||||
file_type = db.Column(db.String(50), nullable=False) # 'json', 'pdf', etc.
|
||||
uploader_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
|
||||
status = db.Column(db.String(20), default='pending') # 'pending', 'processing', 'completed', 'error'
|
||||
description = db.Column(db.Text)
|
||||
uploaded_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
processed_at = db.Column(db.DateTime)
|
||||
|
||||
# ارتباط با آپلودکننده
|
||||
uploader = db.relationship('User', backref='uploaded_files')
|
||||
|
||||
class Article(db.Model):
|
||||
__tablename__ = 'articles'
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
|
||||
article_id = db.Column(db.String(255), nullable=False)
|
||||
name = db.Column(db.String(500), nullable=False)
|
||||
institute = db.Column(db.String(500))
|
||||
abstract = db.Column(db.Text)
|
||||
|
||||
# فیلدهای معیارها
|
||||
measures_positive = db.Column(db.Text)
|
||||
measures_negative = db.Column(db.Text)
|
||||
responses_positive = db.Column(db.Text)
|
||||
responses_negative = db.Column(db.Text)
|
||||
scores_positive = db.Column(db.Text)
|
||||
scores_negative = db.Column(db.Text)
|
||||
|
||||
# فیلدهای جدید برای پنل ادمین
|
||||
pdf_path = db.Column(db.String(500))
|
||||
data_file_id = db.Column(db.Integer, db.ForeignKey('data_files.id'))
|
||||
|
||||
# وضعیت تصحیح
|
||||
correction_status = db.Column(db.String(20), default='pending') # 'pending', 'in_progress', 'completed', 'reviewed'
|
||||
corrected_by = db.Column(db.Integer, db.ForeignKey('users.id'))
|
||||
corrected_at = db.Column(db.DateTime)
|
||||
reviewer_id = db.Column(db.Integer, db.ForeignKey('users.id'))
|
||||
reviewed_at = db.Column(db.DateTime)
|
||||
|
||||
created_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
|
||||
# ارتباطهای موجود - مشخص کردن foreign_keys به صورت صریح
|
||||
user = db.relationship('User', foreign_keys=[user_id], backref=db.backref('articles', lazy=True))
|
||||
|
||||
# ارتباطهای جدید برای پنل ادمین
|
||||
corrector = db.relationship('User', foreign_keys=[corrected_by], backref='corrected_articles')
|
||||
reviewer = db.relationship('User', foreign_keys=[reviewer_id], backref='reviewed_articles')
|
||||
data_file = db.relationship('DataFile', backref='articles')
|
||||
|
||||
class ArticleResponse(db.Model):
|
||||
__tablename__ = 'article_responses'
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
article_id = db.Column(db.Integer, db.ForeignKey('articles.id'), nullable=False)
|
||||
measure_index = db.Column(db.Integer, nullable=False)
|
||||
measure_type = db.Column(db.String(10), nullable=False)
|
||||
original_response = db.Column(db.Text)
|
||||
edited_response = db.Column(db.Text)
|
||||
ai_generated_response = db.Column(db.Text)
|
||||
created_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
|
||||
article = db.relationship('Article', backref=db.backref('responses', lazy=True))
|
||||
|
||||
class ChatSession(db.Model):
|
||||
__tablename__ = 'chat_sessions'
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
|
||||
|
||||
# فیلد جدید برای ارتباط با مقاله
|
||||
article_id = db.Column(db.Integer, db.ForeignKey('articles.id'))
|
||||
|
||||
created_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
|
||||
# ارتباطهای موجود
|
||||
user = db.relationship('User', backref=db.backref('chat_sessions', lazy=True))
|
||||
messages = db.relationship('Message', backref='session', lazy=True, cascade='all, delete-orphan')
|
||||
|
||||
# ارتباط جدید
|
||||
article = db.relationship('Article', backref='chat_sessions')
|
||||
|
||||
class Message(db.Model):
|
||||
__tablename__ = 'messages'
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
session_id = db.Column(db.Integer, db.ForeignKey('chat_sessions.id'), nullable=False)
|
||||
role = db.Column(db.String(20), nullable=False)
|
||||
content = db.Column(db.Text, nullable=False)
|
||||
created_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
|
||||
class CorrectionLog(db.Model):
|
||||
__tablename__ = 'correction_logs'
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
article_id = db.Column(db.Integer, db.ForeignKey('articles.id'), nullable=False)
|
||||
user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
|
||||
action = db.Column(db.String(50), nullable=False) # 'created', 'updated', 'corrected', 'reviewed'
|
||||
field_name = db.Column(db.String(100))
|
||||
old_value = db.Column(db.Text)
|
||||
new_value = db.Column(db.Text)
|
||||
created_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
|
||||
# ارتباطها
|
||||
article = db.relationship('Article', backref='correction_logs')
|
||||
user = db.relationship('User', backref='correction_actions')
|
||||
@@ -0,0 +1,26 @@
|
||||
from utils.config import Config
|
||||
from openai import OpenAI
|
||||
|
||||
|
||||
class Generator:
|
||||
def __init__(self, model_name):
|
||||
super().__init__()
|
||||
self.model_name = model_name
|
||||
self.model = self.config(model_name)
|
||||
|
||||
def generate(self, messages, top_k=None, top_p=None, do_sample=None, max_length=None):
|
||||
print("sending prompt to deepseek")
|
||||
response = self.model.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=messages,
|
||||
stream=False
|
||||
)
|
||||
return response.choices[0].message.content
|
||||
|
||||
def config(self, model_name):
|
||||
return OpenAI(api_key=Config.API_KEY, base_url='https://api.avalai.ir/v1')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
generator = Generator('deepseek-v3.1')
|
||||
print(generator.generate("hi how are you?"))
|
||||
@@ -0,0 +1,92 @@
|
||||
from rag.generator import Generator
|
||||
from rag.retriever import Retriver
|
||||
import numpy as np
|
||||
|
||||
|
||||
class RAG:
|
||||
def __init__(self, model_name, retriever_name, vector_path, data_path, device):
|
||||
self.model_name = model_name
|
||||
self.retriever_name = retriever_name
|
||||
self.vector_path = vector_path
|
||||
self.data_path = data_path
|
||||
self.generator = Generator(model_name)
|
||||
# self.retriever = Retriver(vector_path, data_path, retriever_name, device)
|
||||
|
||||
def make_prompt(self, query, top_k):
|
||||
results = self._retrieve(query, top_k)
|
||||
if results:
|
||||
data = [res.page_content.strip() for res in results]
|
||||
else:
|
||||
data = ["شما محتوایی بارگداری نکرده اید."]
|
||||
prompt = f"""
|
||||
محتوای ارایه شده : {data}
|
||||
سوال : {query}
|
||||
|
||||
لطفا بر اساس محتوای ارایه شده به سوال کامل و دقیق جواب بده دقت کن که بیشتر جواب تولید کنی و دقیق باشی
|
||||
"""
|
||||
sanitized_results = []
|
||||
for res in results:
|
||||
meta = getattr(res, "metadata", {})
|
||||
if isinstance(meta, dict):
|
||||
clean_meta = {}
|
||||
for k, v in meta.items():
|
||||
if isinstance(v, np.generic):
|
||||
clean_meta[k] = v.item()
|
||||
else:
|
||||
clean_meta[k] = v
|
||||
else:
|
||||
clean_meta = {}
|
||||
|
||||
sanitized_doc = {
|
||||
"page_content": res.page_content,
|
||||
"metadata": clean_meta
|
||||
}
|
||||
sanitized_results.append(sanitized_doc)
|
||||
|
||||
return prompt, sanitized_results
|
||||
|
||||
def query(self, query, top_k):
|
||||
results = self._retrieve(query, top_k)
|
||||
if results:
|
||||
data = [res.page_content.strip() for res in results]
|
||||
resp = self.make_response(query, data)
|
||||
else:
|
||||
data = ["شما محتوایی بارگداری نکرده اید."]
|
||||
resp = self.make_response(query, data)
|
||||
|
||||
sanitized_results = []
|
||||
for res in results:
|
||||
meta = getattr(res, "metadata", {})
|
||||
if isinstance(meta, dict):
|
||||
clean_meta = {}
|
||||
for k, v in meta.items():
|
||||
if isinstance(v, np.generic):
|
||||
clean_meta[k] = v.item()
|
||||
else:
|
||||
clean_meta[k] = v
|
||||
else:
|
||||
clean_meta = {}
|
||||
|
||||
sanitized_doc = {
|
||||
"page_content": res.page_content,
|
||||
"metadata": clean_meta
|
||||
}
|
||||
sanitized_results.append(sanitized_doc)
|
||||
|
||||
return sanitized_results, resp
|
||||
|
||||
def make_response(self, query, data):
|
||||
prompt = f"""
|
||||
محتوای ارایه شده : {data}
|
||||
سوال : {query}
|
||||
|
||||
لطفا بر اساس محتوای ارایه شده به سوال کامل و دقیق جواب بده دقت کن که بیشتر جواب تولید کنی و دقیق باشی
|
||||
"""
|
||||
return self.generator.generate(prompt)
|
||||
|
||||
def _retrieve(self, query, top_k):
|
||||
return self.retriever.retrieve(query, top_k)
|
||||
|
||||
def update_retriever(self):
|
||||
"""Retriever را با اسناد جدید در data_path دوباره بارگذاری میکند."""
|
||||
self.retriever.reload_documents()
|
||||
@@ -0,0 +1,106 @@
|
||||
from langchain_community.document_loaders import DirectoryLoader, TextLoader, PyPDFLoader
|
||||
from langchain_huggingface import HuggingFaceEmbeddings
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from langchain_community.vectorstores import FAISS
|
||||
from langchain_core.documents import Document
|
||||
from typing import List, Dict, Any
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
||||
class Retriver:
|
||||
def __init__(self, vector_path, data_path, retriever_name, device):
|
||||
self.vector_path = vector_path
|
||||
self.data_path = data_path
|
||||
self.embedding_model = HuggingFaceEmbeddings(
|
||||
model_name=retriever_name,
|
||||
model_kwargs={"device": device},
|
||||
encode_kwargs={"normalize_embeddings": True}
|
||||
)
|
||||
self.vectorstore = self.load_vectorstore(self.embedding_model, self.vector_path)
|
||||
if self.vectorstore is None:
|
||||
documents = self.load_documents(data_path)
|
||||
chunks = self.split_documents(documents)
|
||||
if chunks:
|
||||
self.vectorstore = self.create_retriever(chunks, self.embedding_model)
|
||||
self.save_vectorstore(self.vectorstore, self.vector_path)
|
||||
|
||||
def get_data_dir(self, data_dir=None):
|
||||
import sys, os
|
||||
if data_dir is not None:
|
||||
return data_dir
|
||||
if getattr(sys, 'frozen', False):
|
||||
exe_dir = os.path.dirname(sys.executable)
|
||||
else:
|
||||
exe_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
data_dir_path = os.path.join(exe_dir, 'data')
|
||||
return data_dir_path
|
||||
|
||||
def load_documents(self, data_dir=None):
|
||||
data_dir = self.get_data_dir(data_dir)
|
||||
pdf_loader = DirectoryLoader(data_dir, glob="*.pdf", loader_cls=PyPDFLoader)
|
||||
txt_loader = DirectoryLoader(data_dir, glob="*.txt", loader_cls=TextLoader)
|
||||
pdf_docs = pdf_loader.load()
|
||||
txt_docs = txt_loader.load()
|
||||
all_docs = pdf_docs + txt_docs
|
||||
return all_docs
|
||||
|
||||
def retrieve(self, query, pdf_path, top_k=10):
|
||||
if not self.vectorstore:
|
||||
return []
|
||||
|
||||
filter_dict = {"source": pdf_path}
|
||||
|
||||
try:
|
||||
docs_and_scores = self.vectorstore.similarity_search_with_score(
|
||||
query,
|
||||
k=top_k,
|
||||
filter=filter_dict
|
||||
)
|
||||
except TypeError:
|
||||
docs_and_scores = self.vectorstore.similarity_search_with_score(query, k=top_k*2)
|
||||
docs_and_scores = [(doc, score) for doc, score in docs_and_scores
|
||||
if doc.metadata.get("source") == pdf_path]
|
||||
|
||||
for doc, score in docs_and_scores:
|
||||
doc.metadata["score"] = float(score)
|
||||
|
||||
return [doc for doc, _ in docs_and_scores[:top_k]]
|
||||
|
||||
def split_documents(self, documents, chunk_size=1000, chunk_overlap=200):
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
separators=["\n\n", "\n", " ", ""]
|
||||
)
|
||||
chunks = text_splitter.split_documents(documents)
|
||||
return chunks
|
||||
|
||||
def create_retriever(self, documents, embedding_model):
|
||||
vectorstore = FAISS.from_documents(documents, embedding_model)
|
||||
return vectorstore
|
||||
|
||||
def save_vectorstore(self, vectorstore, path="./faiss_index"):
|
||||
vectorstore.save_local(path)
|
||||
|
||||
def load_vectorstore(self, embedding_model, path="./faiss_index"):
|
||||
if os.path.exists(os.path.join(path, "index.faiss")):
|
||||
vectorstore = FAISS.load_local(path, embedding_model, allow_dangerous_deserialization=True)
|
||||
return vectorstore
|
||||
else:
|
||||
return None
|
||||
|
||||
def reload_documents(self):
|
||||
"""
|
||||
تمام اسناد موجود در self.data_path را دوباره بارگذاری کرده،
|
||||
vectorstore را بازسازی و ذخیره میکند.
|
||||
"""
|
||||
print("🔄 Reloading documents and rebuilding vectorstore...")
|
||||
documents = self.load_documents(self.data_path)
|
||||
chunks = self.split_documents(documents)
|
||||
if chunks:
|
||||
self.vectorstore = self.create_retriever(chunks, self.embedding_model)
|
||||
self.save_vectorstore(self.vectorstore, self.vector_path)
|
||||
else:
|
||||
self.vectorstore = None
|
||||
print("⚠️ No documents found to reload.")
|
||||
@@ -0,0 +1,119 @@
|
||||
# Python Package Requirements (generated from conda 'nlp' environment)
|
||||
# Note: System/Conda-specific packages (like cuda, blas, compilers) have been filtered out.
|
||||
# It is recommended to use within a Python virtual environment.
|
||||
|
||||
# Core ML & Data Science
|
||||
numpy==1.26.4
|
||||
pandas==2.3.3
|
||||
scipy==1.16.2
|
||||
scikit-learn==1.7.2
|
||||
|
||||
# Deep Learning Frameworks & GPU
|
||||
torch==2.7.1
|
||||
transformers==4.57.1
|
||||
datasets==4.3.0
|
||||
accelerate==1.11.0
|
||||
safetensors==0.6.2
|
||||
bitsandbytes==0.48.1
|
||||
tokenizers==0.22.1
|
||||
faiss-cpu==1.9.0
|
||||
hnswlib==0.8.0
|
||||
|
||||
# LangChain Ecosystem
|
||||
langchain==0.3.27
|
||||
langchain-core==0.3.79
|
||||
langchain-community==0.3.27
|
||||
langchain-openai==0.3.35
|
||||
langchain-anthropic==0.3.22
|
||||
langchain-huggingface==0.3.1
|
||||
langchain-text-splitters==0.3.11
|
||||
langgraph==1.0.1
|
||||
langgraph-checkpoint==3.0.1
|
||||
langgraph-sdk==0.2.9
|
||||
langsmith==0.4.32
|
||||
|
||||
# LLM Libraries & APIs
|
||||
openai==2.8.0
|
||||
anthropic==0.73.0
|
||||
tiktoken==0.12.0
|
||||
huggingface-hub==0.36.0
|
||||
sentence-transformers==5.1.2
|
||||
llama-cpp-python==0.3.16
|
||||
|
||||
# Web & Async
|
||||
fastapi==0.111.0
|
||||
uvicorn==0.29.0
|
||||
starlette==0.37.2
|
||||
pydantic==2.12.3
|
||||
pydantic-core==2.41.4
|
||||
pydantic-settings==2.11.0
|
||||
httpx==0.28.1
|
||||
aiohttp==3.13.1
|
||||
websockets==15.0.1
|
||||
|
||||
# Database
|
||||
sqlalchemy==2.0.44
|
||||
psycopg2-binary==2.9.11
|
||||
pyarrow==22.0.0
|
||||
|
||||
# Flask Ecosystem (if you need it)
|
||||
flask==3.0.3
|
||||
flask-sqlalchemy==3.1.1
|
||||
flask-login==0.6.3
|
||||
flask-bcrypt==1.0.1
|
||||
|
||||
# Utilities & Helpers
|
||||
click==8.3.0
|
||||
colorama==0.4.6
|
||||
loguru==0.7.3
|
||||
tqdm==4.67.1
|
||||
tenacity==9.1.2
|
||||
python-dotenv==1.0.1
|
||||
pyyaml==6.0.3
|
||||
orjson==3.11.4
|
||||
rich==14.2.0
|
||||
typer==0.20.0
|
||||
watchfiles==1.1.1
|
||||
filelock==3.20.0
|
||||
fsspec==2025.9.0
|
||||
jiter==0.12.0
|
||||
sqlite-vec==0.1.6
|
||||
diskcache==5.6.3
|
||||
docstring-parser==0.17.0
|
||||
|
||||
# Jupyter & Notebook (optional)
|
||||
ipython==9.5.0
|
||||
jupyterlab==4.4.7
|
||||
ipykernel==6.30.1
|
||||
traitlets==5.14.3
|
||||
|
||||
# Document Processing
|
||||
pypdf==6.1.3
|
||||
pdfplumber==0.11.5
|
||||
pymupdf==1.24.1
|
||||
python-docx==1.2.0
|
||||
beautifulsoup4==4.13.5
|
||||
lxml==6.0.2
|
||||
|
||||
# Other Dependencies
|
||||
pillow==11.3.0
|
||||
pybind11==3.0.1
|
||||
optree==0.17.0
|
||||
typing-extensions==4.15.0
|
||||
importlib-metadata==8.7.0
|
||||
packaging==25.0
|
||||
requests==2.32.3
|
||||
urllib3==2.5.0
|
||||
certifi==2025.11.12
|
||||
charset-normalizer==3.4.4
|
||||
idna==3.11
|
||||
cryptography==46.0.3
|
||||
bcrypt==5.0.0
|
||||
argon2-cffi==21.3.0
|
||||
python-multipart==0.0.20
|
||||
trustcall==0.0.39
|
||||
dydantic==0.0.8
|
||||
langmem==0.0.30
|
||||
distro==1.9.0
|
||||
webencodings==0.5.1
|
||||
zstandard==0.25.0
|
||||
@@ -0,0 +1,185 @@
|
||||
from fastapi import FastAPI, HTTPException, File, UploadFile
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Dict, Any, Optional
|
||||
import numpy as np
|
||||
import os
|
||||
import shutil
|
||||
|
||||
from agent.espisodic_agent import Agent
|
||||
from utils.config import Config
|
||||
|
||||
|
||||
app = FastAPI(title="YARA RAG API", description="Backend API for YARA RAG assistant")
|
||||
|
||||
agent = Agent(Config)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
os.makedirs(Config.DATA_PATH, exist_ok=True)
|
||||
|
||||
class ChunkMetadata(BaseModel):
|
||||
score: Optional[float] = None
|
||||
source: Optional[str] = None
|
||||
page: Optional[int] = None
|
||||
|
||||
class Chunk(BaseModel):
|
||||
page_content: str
|
||||
metadata: ChunkMetadata
|
||||
|
||||
class QueryRequest(BaseModel):
|
||||
message: str
|
||||
# pdf_path: str
|
||||
|
||||
class QueryResponse(BaseModel):
|
||||
response: str
|
||||
|
||||
|
||||
def sanitize_metadata(metadata: Any) -> Dict[str, Any]:
|
||||
if not isinstance(metadata, dict):
|
||||
return {}
|
||||
|
||||
clean_meta = {}
|
||||
for k, v in metadata.items():
|
||||
if isinstance(v, np.generic):
|
||||
v = v.item()
|
||||
if k == "page":
|
||||
if v is None:
|
||||
clean_meta[k] = None
|
||||
else:
|
||||
try:
|
||||
clean_meta[k] = int(v)
|
||||
except (ValueError, TypeError):
|
||||
clean_meta[k] = None
|
||||
elif k == "score":
|
||||
if v is None:
|
||||
clean_meta[k] = None
|
||||
else:
|
||||
try:
|
||||
clean_meta[k] = float(v)
|
||||
except (ValueError, TypeError):
|
||||
clean_meta[k] = None
|
||||
elif k == "source":
|
||||
clean_meta[k] = str(v) if v is not None else None
|
||||
else:
|
||||
if isinstance(v, (int, float, str, bool, type(None))):
|
||||
clean_meta[k] = v
|
||||
else:
|
||||
clean_meta[k] = str(v)
|
||||
|
||||
return clean_meta
|
||||
|
||||
|
||||
@app.post("/query", response_model=QueryResponse)
|
||||
async def query_rag(request: QueryRequest):
|
||||
if not request.message.strip():
|
||||
raise HTTPException(status_code=400, detail="Message cannot be empty")
|
||||
|
||||
try:
|
||||
print('query callled')
|
||||
full_response = agent.app([{"role": "user", "content": request.message}], 5)
|
||||
print("finisehed")
|
||||
|
||||
return QueryResponse(response=format_response_as_markdown(full_response))
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"RAG processing failed: {str(e)}")
|
||||
|
||||
def format_response_as_markdown(response_text: str) -> str:
|
||||
"""
|
||||
فرمت کردن پاسخ به صورت Markdown برای نمایش زیباتر
|
||||
"""
|
||||
if not response_text:
|
||||
return response_text
|
||||
|
||||
response_text = response_text.strip()
|
||||
|
||||
if any(char in response_text for char in ['#', '*', '`', '-']):
|
||||
return response_text
|
||||
|
||||
paragraphs = [p.strip() for p in response_text.split('\n') if p.strip()]
|
||||
|
||||
formatted_paragraphs = []
|
||||
for i, paragraph in enumerate(paragraphs):
|
||||
if len(paragraph.split()) > 10:
|
||||
formatted_paragraphs.append(paragraph)
|
||||
else:
|
||||
if paragraph.endswith(':') or any(keyword in paragraph.lower() for keyword in ['نکته', 'توجه', 'مهم', 'خلاصه']):
|
||||
formatted_paragraphs.append(f"**{paragraph}**")
|
||||
else:
|
||||
formatted_paragraphs.append(paragraph)
|
||||
|
||||
final_paragraphs = []
|
||||
for paragraph in formatted_paragraphs:
|
||||
if any(indicator in paragraph for indicator in ['•', '- ', '۱.', '۲.', '۳.']):
|
||||
final_paragraphs.append(paragraph)
|
||||
elif len(paragraph.split()) <= 8 and not paragraph.endswith(('.', ':', ';')):
|
||||
final_paragraphs.append(f"- {paragraph}")
|
||||
else:
|
||||
final_paragraphs.append(paragraph)
|
||||
|
||||
return '\n\n'.join(final_paragraphs)
|
||||
|
||||
@app.post("/format_response")
|
||||
async def format_response(request: dict):
|
||||
"""
|
||||
فرمت کردن پاسخ به صورت Markdown
|
||||
"""
|
||||
try:
|
||||
response_text = request.get("response", "")
|
||||
formatted_response = format_response_as_markdown(response_text)
|
||||
|
||||
return {
|
||||
"original": response_text,
|
||||
"formatted": formatted_response,
|
||||
"status": "success"
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
"original": request.get("response", ""),
|
||||
"formatted": request.get("response", ""),
|
||||
"status": "error",
|
||||
"message": str(e)
|
||||
}
|
||||
|
||||
@app.post("/upload")
|
||||
async def upload_files(files: List[UploadFile] = File(...)):
|
||||
if not files:
|
||||
raise HTTPException(status_code=400, detail="No files uploaded")
|
||||
|
||||
saved_files = []
|
||||
try:
|
||||
for file in files:
|
||||
if not file.filename.lower().endswith(".pdf"):
|
||||
raise HTTPException(status_code=400, detail=f"Only PDF files allowed: {file.filename}")
|
||||
|
||||
file_path = os.path.join(Config.DATA_PATH, file.filename)
|
||||
with open(file_path, "wb") as f:
|
||||
shutil.copyfileobj(file.file, f)
|
||||
saved_files.append(file.filename)
|
||||
|
||||
rag.update_retriever()
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"message": f"{len(saved_files)} PDF file(s) uploaded and retriever updated.",
|
||||
"files": saved_files
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Upload or update failed: {str(e)}")
|
||||
|
||||
@app.get("/health")
|
||||
def health_check():
|
||||
return {"status": "OK", "model_loaded": rag is not None}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
@@ -0,0 +1,20 @@
|
||||
from service.rag_api import app as api
|
||||
import uvicorn
|
||||
from werkzeug.serving import make_server
|
||||
|
||||
class Application:
|
||||
def run_fastapi(self):
|
||||
print("FastAPI server starting on http://0.0.0.0:8000")
|
||||
config = uvicorn.Config(
|
||||
api,
|
||||
host="0.0.0.0",
|
||||
port=8000,
|
||||
log_level="info",
|
||||
)
|
||||
self.uvicorn_server = uvicorn.Server(config)
|
||||
self.uvicorn_server.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app = Application()
|
||||
app.run_fastapi()
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# Set Hugging Face mirror for better connectivity
|
||||
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
|
||||
|
||||
class Config:
|
||||
API_KEY = os.getenv("API_KEY")
|
||||
MODEL_NAME = os.getenv("MODEL_NAME")
|
||||
RETREIVER_NAME = os.getenv("RETREIVER_NAME")
|
||||
VECTOR_PATH = os.getenv("VECTOR_PATH")
|
||||
DATA_PATH = os.getenv("DATA_PATH")
|
||||
DEVICE = os.getenv('DEVICE')
|
||||
SECRET_KEY = os.environ.get('SECRET_KEY')
|
||||
SQLALCHEMY_DATABASE_URI = os.environ.get('DATABASE_URL') or 'sqlite:///yara.db'
|
||||
SQLALCHEMY_TRACK_MODIFICATIONS = False
|
||||
RAG_API_URL = os.environ.get('RAG_API_URL') or 'http://localhost:8000/query'
|
||||
SESSION_COOKIE_SECURE = os.environ.get('FLASK_ENV') == 'production'
|
||||
SESSION_COOKIE_HTTPONLY = True
|
||||
SESSION_COOKIE_SAMESITE = 'Lax'
|
||||
SQLALCHEMY_TRACK_MODIFICATIONS = False
|
||||
UPLOAD_FOLDER = 'uploads'
|
||||
# MAX_CONTENT_LENGTH = 50 * 1024 * 1024
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(Config.API_KEY)
|
||||
print(Config.DATA_PATH)
|
||||
print(Config.RETREIVER_NAME)
|
||||
print(Config.DEVICE)
|
||||
Reference in New Issue
Block a user