Create mainapp.py

mainapp.py

@@ -0,0 +1,487 @@
+import os
+import datetime
+import uuid
+import time
+import threading
+import traceback
+import logging
+from queue import Queue # Redisに置き換えるので不要になる
+from dotenv import load_dotenv
+import json
+
+# --- Configuration ---
+load_dotenv()
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+POSTGRES_DSN = os.getenv("POSTGRES_DSN", "postgresql://user:password@localhost:5432/agentdb")
+REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")
+BASE_MODEL_NAME = os.getenv("BASE_MODEL_NAME", "gpt-4o-mini") # Fine-tuning base
+# Fine-tuning するならローカルのOSSモデルが良い場合が多い
+# BASE_MODEL_NAME = "meta-llama/Llama-3-8B-Instruct"
+LEARNING_INTERVAL_HOURS = int(os.getenv("LEARNING_INTERVAL_HOURS", "6"))
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # For PyTorch/TRL
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# --- Library Imports ---
+# (上記 requirements.txt に対応するライブラリを import)
+# LangChain components (as before)
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings # EmbeddingsはHuggingFace製が良いかも
+from langchain.agents import AgentExecutor, create_react_agent, Tool
+# ... other langchain imports
+
+# Database (SQLAlchemy example)
+from sqlalchemy import create_engine, Column, Integer, String, Float, Boolean, DateTime, Text, MetaData, Index
+from sqlalchemy.dialects.postgresql import UUID, JSONB # Use BYTEA or pgvector extension for vectors
+# from sqlalchemy.dialects.postgresql import BYTEA # For raw byte vectors
+# from pgvector.sqlalchemy import Vector # If using pgvector extension
+from sqlalchemy.orm import sessionmaker, declarative_base
+import sqlalchemy # Ensure it's imported
+
+# Message Queue
+import redis
+
+# Vectorization
+from sentence_transformers import SentenceTransformer
+
+# Scheduling
+from apscheduler.schedulers.background import BackgroundScheduler
+from apscheduler.triggers.interval import IntervalTrigger
+
+# TRL (Placeholders for actual imports and usage)
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import LoraConfig
+from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead, create_reference_model
+from trl.core import LengthSampler
+
+# --- Database Setup (SQLAlchemy) ---
+Base = declarative_base()
+engine = create_engine(POSTGRES_DSN)
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+# Example Experience Table (Needs pgvector extension or BYTEA for vectors)
+class Experience(Base):
+    __tablename__ = "experiences"
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    timestamp = Column(DateTime, default=datetime.datetime.utcnow)
+    goal = Column(Text)
+    task = Column(Text)
+    # thought_summary = Column(Text) # Storing full thoughts can be large
+    action_info = Column(JSONB) # Store action, input, tool used etc.
+    observation_summary = Column(Text) # Summarize or store key parts
+    success = Column(Boolean)
+    feedback_score = Column(Float, default=0.0) # Numerical feedback
+    execution_time = Column(Float)
+    # --- Vector Representations ---
+    # Option 1: Use pgvector extension (Recommended)
+    # task_vector = Column(Vector(384)) # Example dimension for all-MiniLM-L6-v2
+    # observation_vector = Column(Vector(384))
+    # state_vector = Column(Vector(768)) # Example combined vector
+    # __table_args__ = (Index('ix_experiences_state_vector', state_vector, postgresql_using='hnsw', postgresql_with={'m': 16, 'ef_construction': 64}),)
+
+    # Option 2: Use BYTEA (Requires manual handling of bytes)
+    # task_vector_bytes = Column(BYTEA)
+    # observation_vector_bytes = Column(BYTEA)
+
+# Example Task Table
+class Task(Base):
+    __tablename__ = "tasks"
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    goal = Column(Text)
+    task_description = Column(Text)
+    status = Column(String, default="pending") # pending, processing, completed, failed
+    created_at = Column(DateTime, default=datetime.datetime.utcnow)
+    updated_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow)
+    result = Column(Text, nullable=True)
+
+# Create tables if they don't exist
+Base.metadata.create_all(bind=engine)
+
+# --- Message Queue Setup (Redis) ---
+redis_client = redis.from_url(REDIS_URL, decode_responses=True)
+TASK_QUEUE_KEY = "agent_task_queue"
+
+# --- Vectorization Model ---
+# Use a sentence transformer model suitable for tasks/observations
+# Consider models optimized for semantic similarity.
+# Run this on CPU or GPU depending on availability/need.
+embedding_model_name = 'all-MiniLM-L6-v2' # Example model
+logging.info(f"Loading sentence transformer model: {embedding_model_name}...")
+# Specify device to control CPU/GPU usage for embeddings
+sentence_model = SentenceTransformer(embedding_model_name, device='cpu') # Use CPU for potentially less conflict with TRL on GPU
+logging.info("Sentence transformer model loaded.")
+
+def get_vector(text: str):
+    """Generates a vector embedding for the given text."""
+    if not text:
+        return None
+    # Ensure model is on the correct device if moved
+    # sentence_model.to('cpu')
+    vector = sentence_model.encode(text, convert_to_numpy=True)
+    # If using BYTEA: return vector.tobytes()
+    # If using pgvector: return vector.tolist() # Or directly numpy array if supported
+    return vector.tolist() # For pgvector
+
+# --- Experience Management (using DB) ---
+def add_experience_db(task_info: dict, agent_output: dict, success: bool, feedback: float = 0.0, exec_time: float = 0.0):
+    """Adds an agent's experience to the PostgreSQL database."""
+    db = SessionLocal()
+    try:
+        # --- Generate Vector Representations ---
+        task_vector = get_vector(task_info.get("task"))
+        obs_summary = agent_output.get("output", "")[:500] # Limit observation size
+        observation_vector = get_vector(obs_summary)
+        # Combine vectors or create a more complex state representation
+        state_vector = None
+        if task_vector and observation_vector:
+             # Simple concatenation example (ensure dimensions match DB schema)
+             # state_vector = task_vector + observation_vector
+             pass # Implement actual state vector logic
+
+        action_info = {
+            "action": agent_output.get("action", "unknown"), # Extract action if available
+            "input": agent_output.get("action_input", "unknown"), # Extract input if available
+            # Add other relevant details like tool used
+        }
+
+        exp = Experience(
+            goal=task_info.get("goal"),
+            task=task_info.get("task"),
+            action_info=action_info,
+            observation_summary=obs_summary,
+            success=success,
+            feedback_score=feedback,
+            execution_time=exec_time,
+            # task_vector=task_vector, # Assign vectors (match DB column type)
+            # observation_vector=observation_vector,
+            # state_vector=state_vector,
+        )
+        db.add(exp)
+        db.commit()
+        logging.debug(f"Experience added to DB: Success={success}, Task={task_info.get('task')[:50]}")
+    except Exception as e:
+        db.rollback()
+        logging.error(f"Failed to add experience to DB: {e}", exc_info=True)
+    finally:
+        db.close()
+
+def retrieve_relevant_experiences_db(query: str, k: int = 3) -> list[Experience]:
+    """Retrieves relevant experiences using vector similarity search (requires pgvector)."""
+    db = SessionLocal()
+    try:
+        query_vector = get_vector(query)
+        if query_vector is None:
+            return []
+
+        # --- Requires pgvector setup ---
+        # This query syntax depends on sqlalchemy-pgvector or raw SQL
+        # results = db.query(Experience).order_by(Experience.state_vector.l2_distance(query_vector)).limit(k).all()
+        # logging.info(f"Retrieved {len(results)} experiences from DB for query: {query[:50]}")
+        # return results
+
+        # --- Placeholder if pgvector is not set up ---
+        logging.warning("Vector search in DB requested but not implemented (requires pgvector). Returning empty list.")
+        return []
+    except Exception as e:
+        logging.error(f"Failed to retrieve experiences from DB: {e}", exc_info=True)
+        return []
+    finally:
+        db.close()
+
+# --- Tools Definition (same as before) ---
+# ... search, python_repl ...
+tools = [
+    Tool(name="Search", func=search.run, description="..."),
+    Tool(name="PythonREPL", func=python_repl.run, description="..."),
+]
+
+# --- Agent Setup ---
+# Use the base model for the agent initially. The fine-tuned model will be loaded by the learning worker.
+agent_llm = ChatOpenAI(model=BASE_MODEL_NAME, temperature=0.3, api_key=OPENAI_API_KEY)
+prompt_template = hub.pull("hwchase17/react-chat")
+agent = create_react_agent(agent_llm, tools, prompt_template)
+agent_executor = AgentExecutor(
+    agent=agent, tools=tools, verbose=False, handle_parsing_errors=True, max_iterations=10,
+)
+
+
+# --- Learning Module (TRL Implementation Sketch) ---
+learning_lock = threading.Lock()
+ppo_trainer = None # Global PPO trainer instance (or manage per learning cycle)
+fine_tuned_model_path = "./fine_tuned_model" # Path to save/load fine-tuned adapter/model
+
+def calculate_reward(experience_data: dict) -> float:
+    """Calculates a reward score based on experience."""
+    reward = 0.0
+    if experience_data.get("success"):
+        reward += 1.0
+    else:
+        reward -= 1.0 # Penalty for failure
+
+    # Penalty for long execution time (log scale to moderate impact)
+    exec_time = experience_data.get("execution_time", 1.0) # Avoid log(0)
+    if exec_time > 1.0:
+        reward -= 0.1 * min(max(0, exec_time), 300)**0.5 # Capped sqrt penalty
+
+    # Incorporate feedback score
+    reward += experience_data.get("feedback_score", 0.0) * 0.5 # Scale feedback impact
+
+    return reward
+
+def prepare_ppo_data(experiences: list[Experience]) -> list[dict]:
+    """Prepares data in the format expected by TRL's PPOTrainer."""
+    ppo_data = []
+    for exp in experiences:
+        # Construct the 'query' - the input to the LLM for the task
+        query_text = f"Goal: {exp.goal}\nTask: {exp.task}"
+        # Construct the 'response' - the LLM's actual output (observation)
+        response_text = exp.observation_summary
+        # Calculate reward
+        reward_score = calculate_reward(exp.metadata) # Assuming metadata is attached or retrieved
+
+        if query_text and response_text:
+             ppo_data.append({
+                 "query": query_text,
+                 "response": response_text,
+                 "reward": torch.tensor([reward_score], dtype=torch.float3_tensors) # TRL expects tensor
+             })
+    return ppo_data
+
+
+def run_learning_cycle():
+    """The main learning process using TRL."""
+    global ppo_trainer # Allow modification
+    if not torch.cuda.is_available():
+        logging.warning("CUDA not available. Skipping fine-tuning cycle.")
+        return
+
+    with learning_lock:
+        logging.info(f"[Learning Cycle Triggered] - Device: {DEVICE}")
+        start_time = time.time()
+
+        # 1. Fetch Data from PostgreSQL
+        logging.info("Fetching recent experiences from PostgreSQL...")
+        db = SessionLocal()
+        try:
+            # Fetch experiences (e.g., last N or within a time window)
+            recent_experiences = db.query(Experience).order_by(Experience.timestamp.desc()).limit(500).all() # Adjust limit
+        finally:
+            db.close()
+
+        if not recent_experiences or len(recent_experiences) < 50: # Need sufficient data
+            logging.info(f"Not enough new experiences ({len(recent_experiences)}). Skipping fine-tuning.")
+            return
+        logging.info(f"Fetched {len(recent_experiences)} experiences for learning.")
+
+        # 2. Prepare Data and Calculate Rewards
+        logging.info("Preparing data for PPO...")
+        ppo_data = prepare_ppo_data(recent_experiences)
+        if not ppo_data:
+             logging.warning("No valid data points after preparation. Skipping fine-tuning.")
+             return
+
+        # Convert to TRL dataset format (example, check TRL docs for specifics)
+        # This usually involves tokenizing queries and responses
+        # query_tensors = [tokenizer.encode(d['query'], return_tensors="pt").squeeze(0) for d in ppo_data]
+        # response_tensors = [tokenizer.encode(d['response'], return_tensors="pt").squeeze(0) for d in ppo_data]
+        # rewards = [d['reward'] for d in ppo_data]
+
+        # 3. Setup TRL PPO Trainer (Simplified Example)
+        logging.info("Setting up TRL PPOTrainer...")
+        try:
+            # --- TRL Configuration ---
+            ppo_config = PPOConfig(
+                model_name=BASE_MODEL_NAME,
+                learning_rate=1.41e-5,
+                batch_size=16, # Adjust based on GPU memory
+                mini_batch_size=4, # Adjust based on GPU memory
+                gradient_accumulation_steps=1,
+                optimize_cuda_cache=True,
+                # early_stopping=True,
+                # target_kl=0.1,
+                ppo_epochs=4, # Number of epochs per PPO step
+                seed=42,
+                # Use LoRA for efficient fine-tuning
+                use_lora=True,
+            )
+
+            # --- Model Loading (with Quantization and LoRA) ---
+            # bnb_config = BitsAndBytesConfig(...) # Optional quantization
+            lora_config = LoraConfig(
+                 r=16, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
+            )
+            tokenizer = AutoTokenizer.from_pretrained(ppo_config.model_name)
+            if getattr(tokenizer, "pad_token", None) is None:
+                 tokenizer.pad_token = tokenizer.eos_token # Important for padding
+
+            # Load the base model with ValueHead for PPO and LoRA config
+            model = AutoModelForCausalLMWithValueHead.from_pretrained(
+                ppo_config.model_name,
+                # quantization_config=bnb_config, # Optional
+                peft_config=lora_config,
+                # load_in_8bit=True, # Or load_in_4bit=True
+                torch_dtype=torch.float16, # Use float16/bfloat16 on GPU
+                device_map="auto" # Use Accelerate for device mapping
+            )
+            # Reference model for KL divergence
+            ref_model = create_reference_model(model) # Or load separately
+
+            # --- Initialize Trainer ---
+            # Requires tokenized queries, responses, and rewards
+            # ppo_trainer = PPOTrainer(
+            #     config=ppo_config,
+            #     model=model,
+            #     ref_model=ref_model,
+            #     tokenizer=tokenizer,
+            #     dataset=your_prepared_dataset, # Requires tokenized data
+            #     data_collator=your_data_collator # Handles padding
+            # )
+
+            # --- PPO Training Loop ---
+            logging.info("Starting PPO Training Loop (Simulation - Actual requires dataset)...")
+            # for epoch in range(ppo_config.ppo_epochs):
+            #     for batch in ppo_trainer.dataloader:
+            #         # Get query tensors, response tensors from batch
+            #         # Compute log probs, values, etc.
+            #         # stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
+            #         # ppo_trainer.log_stats(stats, batch, rewards)
+            #     # Save model checkpoint periodically?
+            time.sleep(10) # Simulate training time
+
+            # --- Save Fine-tuned Model (LoRA Adapters) ---
+            logging.info("Saving fine-tuned LoRA adapters...")
+            # ppo_trainer.save_pretrained(fine_tuned_model_path)
+            # tokenizer.save_pretrained(fine_tuned_model_path)
+            logging.info(f"Fine-tuned adapters saved to {fine_tuned_model_path}")
+
+        except Exception as e:
+            logging.error(f"Error during TRL setup or training: {e}", exc_info=True)
+            # Clean up GPU memory if needed
+            del model, ref_model, ppo_trainer
+            torch.cuda.empty_cache()
+
+        logging.info(f"Learning cycle finished. Duration: {time.time() - start_time:.2f}s")
+
+# --- Task Management (using Redis) ---
+def add_task_mq(task: str, goal: str):
+    """Adds a task to the Redis queue."""
+    task_id = str(uuid.uuid4())
+    task_data = json.dumps({"id": task_id, "task": task, "goal": goal})
+    try:
+        redis_client.lpush(TASK_QUEUE_KEY, task_data)
+        logging.info(f"Task {task_id} added to Redis queue: {task[:50]}...")
+    except Exception as e:
+        logging.error(f"Failed to add task to Redis: {e}")
+
+# --- Agent Worker (modified for Redis and DB) ---
+def agent_worker(worker_id: int):
+    """Processes tasks from the Redis queue."""
+    logging.info(f"Agent Worker-{worker_id} started.")
+    while True: # Run continuously
+        try:
+            # Blocking pop from Redis list (wait indefinitely)
+            _, task_data_json = redis_client.brpop(TASK_QUEUE_KEY)
+            task_info = json.loads(task_data_json)
+            task_id = task_info["id"]
+            task_desc = task_info["task"]
+            goal = task_info["goal"]
+
+            logging.info(f"Worker-{worker_id} processing Task {task_id}: {task_desc[:50]}...")
+            start_time = time.time()
+            success = False
+            final_output = None
+            agent_result = {} # Store agent's output details
+
+            # Update task status in DB (optional)
+            # update_task_status(task_id, "processing")
+
+            # --- Retrieve relevant experiences ---
+            # query = f"Goal: {goal}\nTask: {task_desc}"
+            # relevant_experiences = retrieve_relevant_experiences_db(query, k=3)
+            # experience_context = ... # Format context from DB results
+
+            # --- Prepare Agent Input ---
+            input_messages = [
+                SystemMessage(content=f"Your long term goal is: {goal}. Think step-by-step."),
+                # Add experience_context here if needed
+                HumanMessage(content=f"Current task: {task_desc}")
+            ]
+
+            # --- Execute Agent ---
+            try:
+                # Ideally, load the latest fine-tuned model for inference here
+                # This requires coordination or loading the adapter weights
+                agent_result = agent_executor.invoke({"input": input_messages})
+                final_output = agent_result.get("output", "No output.")
+                # Simple success check (refine this based on tool usage, keywords etc.)
+                success = not any(err in final_output.lower() for err in ["error", "fail", "unable"])
+            except Exception as e:
+                logging.error(f"Worker-{worker_id} Task {task_id} failed during execution: {e}", exc_info=True)
+                final_output = f"Agent execution failed: {e}"
+                success = False
+                agent_result = {"output": final_output, "action": "error"} # Log error state
+
+            # --- Record Experience ---
+            exec_time = time.time() - start_time
+            # Add user feedback later if available (e.g., via API)
+            feedback_score = 0.0
+            add_experience_db(task_info, agent_result, success, feedback_score, exec_time)
+
+            # Update task status in DB (optional)
+            # update_task_status(task_id, "completed" if success else "failed", final_output)
+
+            logging.info(f"Worker-{worker_id} finished Task {task_id}. Success: {success}. Time: {exec_time:.2f}s")
+
+        except redis.exceptions.ConnectionError as e:
+             logging.error(f"Worker-{worker_id} Redis connection error: {e}. Retrying in 10s...")
+             time.sleep(10)
+        except Exception as e:
+            logging.error(f"Worker-{worker_id} encountered an unexpected error: {e}", exc_info=True)
+            time.sleep(5) # Avoid rapid looping on persistent errors
+
+# --- Main Execution / Service Startup ---
+if __name__ == "__main__":
+    logging.info("Initializing Agent System...")
+
+    # --- Start Background Learning Scheduler ---
+    scheduler = BackgroundScheduler(daemon=True)
+    scheduler.add_job(
+        run_learning_cycle,
+        trigger=IntervalTrigger(hours=LEARNING_INTERVAL_HOURS),
+        id="learning_job",
+        name="Fine-tuning Learning Cycle",
+        replace_existing=True
+    )
+    scheduler.start()
+    logging.info(f"Background learning scheduler started. Interval: {LEARNING_INTERVAL_HOURS} hours.")
+
+    # --- Start Agent Workers ---
+    num_workers = int(os.getenv("NUM_WORKERS", "2"))
+    worker_threads = []
+    for i in range(num_workers):
+        thread = threading.Thread(target=agent_worker, args=(i+1,), daemon=True)
+        thread.start()
+        worker_threads.append(thread)
+    logging.info(f"{num_workers} Agent worker threads started.")
+
+    # --- Add Initial Tasks (Example) ---
+    add_task_mq("Explain the difference between LoRA and full fine-tuning for LLMs.",
+                "Understand AI model optimization techniques.")
+    add_task_mq("Write a Python script using pandas to read a CSV file named 'data.csv' and print the first 5 rows.",
+                "Develop data processing scripts.")
+
+    logging.info("Agent system is running. Workers processing tasks from Redis.")
+    logging.info("Press Ctrl+C to stop.")
+
+    try:
+        # Keep main thread alive
+        while True:
+            time.sleep(60)
+            # Add health checks or monitoring here if needed
+            logging.debug("Main thread alive...")
+    except KeyboardInterrupt:
+        logging.info("Shutdown signal received...")
+        scheduler.shutdown()
+        # Workers are daemon threads, they will exit when main thread exits.
+        # Implement graceful shutdown for workers if needed (e.g., sending sentinel)
+        logging.info("Agent system stopped.")