groknroll 2.0.0__py3-none-any.whl

groknroll/storage/database.py

@@ -0,0 +1,272 @@
+"""
+Database management for groknroll
+"""
+
+from pathlib import Path
+from typing import Optional, List, Dict, Any
+from datetime import datetime
+
+from sqlalchemy import create_engine, desc
+from sqlalchemy.orm import sessionmaker, Session as DBSession
+
+from groknroll.storage.models import Base, Project, FileIndex, Execution, Session, Analysis
+
+
+class Database:
+    """
+    SQLite database for groknroll project state
+
+    Stores:
+    - Project metadata
+    - File index with AST data
+    - RLM execution history
+    - Session history
+    - Analysis results
+    """
+
+    def __init__(self, db_path: Optional[Path] = None):
+        """
+        Initialize database
+
+        Args:
+            db_path: Path to SQLite database file.
+                    If None, uses ~/.groknroll/groknroll.db
+        """
+        if db_path is None:
+            db_path = Path.home() / ".groknroll" / "groknroll.db"
+
+        db_path.parent.mkdir(parents=True, exist_ok=True)
+
+        self.db_path = db_path
+        self.engine = create_engine(f"sqlite:///{db_path}", echo=False)
+        self.SessionLocal = sessionmaker(bind=self.engine)
+
+        # Create tables
+        Base.metadata.create_all(self.engine)
+
+    def get_session(self) -> DBSession:
+        """Get database session"""
+        return self.SessionLocal()
+
+    # =========================================================================
+    # Project Operations
+    # =========================================================================
+
+    def get_or_create_project(self, project_path: Path) -> Project:
+        """Get existing project or create new one"""
+        with self.get_session() as session:
+            project = session.query(Project).filter_by(path=str(project_path)).first()
+
+            if project is None:
+                project = Project(
+                    path=str(project_path),
+                    name=project_path.name
+                )
+                session.add(project)
+                session.commit()
+                session.refresh(project)
+
+            return project
+
+    def update_project_stats(
+        self,
+        project_id: int,
+        total_files: int,
+        total_lines: int
+    ) -> None:
+        """Update project statistics"""
+        with self.get_session() as session:
+            project = session.query(Project).filter_by(id=project_id).first()
+            if project:
+                project.total_files = total_files
+                project.total_lines = total_lines
+                project.last_indexed = datetime.utcnow()
+                session.commit()
+
+    # =========================================================================
+    # File Index Operations
+    # =========================================================================
+
+    def index_file(
+        self,
+        project_id: int,
+        file_path: Path,
+        relative_path: str,
+        **metadata
+    ) -> FileIndex:
+        """Index a file"""
+        with self.get_session() as session:
+            file_index = session.query(FileIndex).filter_by(
+                project_id=project_id,
+                path=str(file_path)
+            ).first()
+
+            if file_index is None:
+                file_index = FileIndex(
+                    project_id=project_id,
+                    path=str(file_path),
+                    relative_path=relative_path
+                )
+                session.add(file_index)
+
+            # Update metadata
+            for key, value in metadata.items():
+                if hasattr(file_index, key):
+                    setattr(file_index, key, value)
+
+            file_index.updated_at = datetime.utcnow()
+            session.commit()
+            session.refresh(file_index)
+
+            return file_index
+
+    def get_project_files(self, project_id: int) -> List[FileIndex]:
+        """Get all indexed files for project"""
+        with self.get_session() as session:
+            return session.query(FileIndex).filter_by(project_id=project_id).all()
+
+    def search_files(
+        self,
+        project_id: int,
+        query: str,
+        language: Optional[str] = None
+    ) -> List[FileIndex]:
+        """Search indexed files"""
+        with self.get_session() as session:
+            q = session.query(FileIndex).filter_by(project_id=project_id)
+
+            if language:
+                q = q.filter_by(language=language)
+
+            q = q.filter(FileIndex.relative_path.like(f"%{query}%"))
+
+            return q.all()
+
+    # =========================================================================
+    # Execution History
+    # =========================================================================
+
+    def log_execution(
+        self,
+        project_id: int,
+        task: str,
+        response: str,
+        **metrics
+    ) -> Execution:
+        """Log RLM execution"""
+        with self.get_session() as session:
+            execution = Execution(
+                project_id=project_id,
+                task=task,
+                response=response,
+                status=metrics.get("status", "success"),
+                **metrics
+            )
+            session.add(execution)
+            session.commit()
+            session.refresh(execution)
+
+            return execution
+
+    def get_recent_executions(
+        self,
+        project_id: int,
+        limit: int = 10
+    ) -> List[Execution]:
+        """Get recent executions"""
+        with self.get_session() as session:
+            return (
+                session.query(Execution)
+                .filter_by(project_id=project_id)
+                .order_by(desc(Execution.started_at))
+                .limit(limit)
+                .all()
+            )
+
+    def get_execution_stats(self, project_id: int) -> Dict[str, Any]:
+        """Get execution statistics"""
+        with self.get_session() as session:
+            executions = session.query(Execution).filter_by(project_id=project_id).all()
+
+            total_cost = sum(e.total_cost or 0.0 for e in executions)
+            total_time = sum(e.total_time or 0.0 for e in executions)
+            total_count = len(executions)
+            success_count = sum(1 for e in executions if e.status == "success")
+
+            return {
+                "total_executions": total_count,
+                "successful": success_count,
+                "failed": total_count - success_count,
+                "total_cost": total_cost,
+                "total_time": total_time,
+                "avg_cost": total_cost / total_count if total_count > 0 else 0.0,
+                "avg_time": total_time / total_count if total_count > 0 else 0.0,
+            }
+
+    # =========================================================================
+    # Session Management
+    # =========================================================================
+
+    def create_session(
+        self,
+        project_id: int,
+        session_type: str
+    ) -> Session:
+        """Create new session"""
+        with self.get_session() as session:
+            new_session = Session(
+                project_id=project_id,
+                session_type=session_type
+            )
+            session.add(new_session)
+            session.commit()
+            session.refresh(new_session)
+
+            return new_session
+
+    def end_session(self, session_id: int) -> None:
+        """End session"""
+        with self.get_session() as session:
+            sess = session.query(Session).filter_by(id=session_id).first()
+            if sess:
+                sess.ended_at = datetime.utcnow()
+                session.commit()
+
+    # =========================================================================
+    # Analysis Results
+    # =========================================================================
+
+    def save_analysis(
+        self,
+        project_id: int,
+        analysis_type: str,
+        results: Dict[str, Any],
+        **metadata
+    ) -> Analysis:
+        """Save analysis results"""
+        with self.get_session() as session:
+            analysis = Analysis(
+                project_id=project_id,
+                analysis_type=analysis_type,
+                results=results,
+                **metadata
+            )
+            session.add(analysis)
+            session.commit()
+            session.refresh(analysis)
+
+            return analysis
+
+    def get_latest_analysis(
+        self,
+        project_id: int,
+        analysis_type: str
+    ) -> Optional[Analysis]:
+        """Get latest analysis of given type"""
+        with self.get_session() as session:
+            return (
+                session.query(Analysis)
+                .filter_by(project_id=project_id, analysis_type=analysis_type)
+                .order_by(desc(Analysis.created_at))
+                .first()
+            )

groknroll/storage/models.py

@@ -0,0 +1,128 @@
+"""
+Database models for groknroll
+"""
+
+from datetime import datetime
+from typing import Optional
+from sqlalchemy import Column, Integer, String, Float, DateTime, Text, ForeignKey, Boolean, JSON
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import relationship
+
+Base = declarative_base()
+
+
+class Project(Base):
+    """Project metadata"""
+    __tablename__ = "projects"
+
+    id = Column(Integer, primary_key=True)
+    path = Column(String, unique=True, nullable=False, index=True)
+    name = Column(String, nullable=False)
+    language = Column(String)  # Primary language
+    total_files = Column(Integer, default=0)
+    total_lines = Column(Integer, default=0)
+    last_indexed = Column(DateTime)
+    created_at = Column(DateTime, default=datetime.utcnow)
+    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+
+    # Relationships
+    files = relationship("FileIndex", back_populates="project", cascade="all, delete-orphan")
+    executions = relationship("Execution", back_populates="project", cascade="all, delete-orphan")
+    sessions = relationship("Session", back_populates="project", cascade="all, delete-orphan")
+    analyses = relationship("Analysis", back_populates="project", cascade="all, delete-orphan")
+
+
+class FileIndex(Base):
+    """Indexed file metadata"""
+    __tablename__ = "file_index"
+
+    id = Column(Integer, primary_key=True)
+    project_id = Column(Integer, ForeignKey("projects.id"), nullable=False, index=True)
+    path = Column(String, nullable=False, index=True)
+    relative_path = Column(String, nullable=False)
+    language = Column(String)
+    size_bytes = Column(Integer)
+    lines_of_code = Column(Integer)
+    complexity = Column(Float)  # Cyclomatic complexity
+    last_modified = Column(DateTime)
+    ast_data = Column(JSON)  # Parsed AST metadata
+    imports = Column(JSON)  # List of imports
+    exports = Column(JSON)  # List of exports
+    created_at = Column(DateTime, default=datetime.utcnow)
+    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+
+    # Relationships
+    project = relationship("Project", back_populates="files")
+
+
+class Execution(Base):
+    """RLM execution history"""
+    __tablename__ = "executions"
+
+    id = Column(Integer, primary_key=True)
+    project_id = Column(Integer, ForeignKey("projects.id"), index=True)
+    session_id = Column(Integer, ForeignKey("sessions.id"), index=True)
+
+    # Request
+    task = Column(Text, nullable=False)
+    context = Column(JSON)
+    model = Column(String)
+
+    # Response
+    response = Column(Text)
+    trace_log = Column(Text)
+
+    # Metrics
+    total_cost = Column(Float)
+    total_time = Column(Float)
+    iterations = Column(Integer)
+    status = Column(String)  # success, failed, timeout
+    error_message = Column(Text)
+
+    # Timestamps
+    started_at = Column(DateTime, default=datetime.utcnow, index=True)
+    completed_at = Column(DateTime)
+    created_at = Column(DateTime, default=datetime.utcnow)
+
+    # Relationships
+    project = relationship("Project", back_populates="executions")
+    session = relationship("Session", back_populates="executions")
+
+
+class Session(Base):
+    """Interactive session history"""
+    __tablename__ = "sessions"
+
+    id = Column(Integer, primary_key=True)
+    project_id = Column(Integer, ForeignKey("projects.id"), index=True)
+
+    session_type = Column(String)  # chat, repl, dashboard
+    started_at = Column(DateTime, default=datetime.utcnow, index=True)
+    ended_at = Column(DateTime)
+    message_count = Column(Integer, default=0)
+    total_cost = Column(Float, default=0.0)
+
+    # Relationships
+    project = relationship("Project", back_populates="sessions")
+    executions = relationship("Execution", back_populates="session", cascade="all, delete-orphan")
+
+
+class Analysis(Base):
+    """Code analysis results"""
+    __tablename__ = "analyses"
+
+    id = Column(Integer, primary_key=True)
+    project_id = Column(Integer, ForeignKey("projects.id"), index=True)
+
+    analysis_type = Column(String, nullable=False)  # security, complexity, review, etc
+    target_path = Column(String)  # File or directory analyzed
+    results = Column(JSON)  # Analysis results
+    recommendations = Column(JSON)  # Recommendations
+    issues = Column(JSON)  # Issues found
+    metrics = Column(JSON)  # Metrics
+
+    execution_time = Column(Float)
+    created_at = Column(DateTime, default=datetime.utcnow, index=True)
+
+    # Relationships
+    project = relationship("Project", back_populates="analyses")

groknroll/utils/parsing.py

@@ -0,0 +1,168 @@
+"""
+Parsing utilities for RLM trajectories.
+"""
+
+import re
+from typing import TYPE_CHECKING
+
+from groknroll.core.types import REPLResult, RLMIteration
+
+if TYPE_CHECKING:
+    from groknroll.environments.base_env import BaseEnv
+
+
+def find_code_blocks(text: str) -> list[str]:
+    """
+    Find REPL code blocks in text wrapped in triple backticks and return List of content(s).
+    Returns None if no code blocks are found.
+    """
+    pattern = r"```repl\s*\n(.*?)\n```"
+    results = []
+
+    for match in re.finditer(pattern, text, re.DOTALL):
+        code_content = match.group(1).strip()
+        results.append(code_content)
+
+    return results
+
+
+def find_final_answer(text: str, environment: "BaseEnv | None" = None) -> str | None:
+    """
+    Find FINAL(...) or FINAL_VAR(...) statement in response and return the final answer string.
+
+    If FINAL_VAR is found and an environment is provided, executes code to retrieve the variable value.
+    Returns None if neither pattern is found.
+
+    Args:
+        text: The response text to parse
+        environment: Optional environment to execute code for FINAL_VAR retrieval
+
+    Returns:
+        The final answer string, or None if no final answer pattern is found
+    """
+    # Check for FINAL_VAR pattern first - must be at start of line
+    final_var_pattern = r"^\s*FINAL_VAR\((.*?)\)"
+    match = re.search(final_var_pattern, text, re.MULTILINE | re.DOTALL)
+    if match:
+        variable_name = match.group(1).strip().strip('"').strip("'")
+        if environment is not None:
+            result = environment.execute_code(f"print(FINAL_VAR({variable_name!r}))")
+            final_answer = result.stdout.strip()
+            if final_answer == "":
+                final_answer = result.stderr.strip() or ""
+            return final_answer
+        return None
+
+    # Check for FINAL pattern - must be at start of line
+    final_pattern = r"^\s*FINAL\((.*?)\)"
+    match = re.search(final_pattern, text, re.MULTILINE | re.DOTALL)
+    if match:
+        return match.group(1).strip()
+
+    return None
+
+
+def format_iteration(
+    iteration: RLMIteration, max_character_length: int = 20000
+) -> list[dict[str, str]]:
+    """
+    Format an RLM iteration (including all code blocks) to append to the message history for
+    the prompt of the LM in the next iteration. We also truncate code execution results
+    that exceed the max_character_length.
+
+    Args:
+        iteration: The iteration to format
+        max_character_length: The maximum character length of the result
+
+    Returns:
+        A list of messages to add to the next prompt
+    """
+    messages = [{"role": "assistant", "content": iteration.response}]
+
+    for code_block in iteration.code_blocks:
+        code = code_block.code
+        result = code_block.result
+        result = format_execution_result(result)
+        if len(result) > max_character_length:
+            result = (
+                result[:max_character_length]
+                + f"... + [{len(result) - max_character_length} chars...]"
+            )
+
+        execution_message = {
+            "role": "user",
+            "content": f"Code executed:\n```python\n{code}\n```\n\nREPL output:\n{result}",
+        }
+        messages.append(execution_message)
+    return messages
+
+
+################
+# TODO: Remove and refactor these soon
+################
+
+
+def format_execution_result(result: REPLResult) -> str:
+    """
+    Format the execution result as a string for display.
+
+    Args:
+        result: The REPLResult object to format.
+    """
+    result_parts = []
+
+    if result.stdout:
+        result_parts.append(f"\n{result.stdout}")
+
+    if result.stderr:
+        result_parts.append(f"\n{result.stderr}")
+
+    # Show some key variables (excluding internal ones)
+    important_vars = {}
+    for key, value in result.locals.items():
+        if not key.startswith("_") and key not in [
+            "__builtins__",
+            "__name__",
+            "__doc__",
+        ]:
+            # Only show simple types or short representations
+            if isinstance(value, (str, int, float, bool, list, dict, tuple)):
+                important_vars[key] = ""
+
+    if important_vars:
+        result_parts.append(f"REPL variables: {list(important_vars.keys())}\n")
+
+    return "\n\n".join(result_parts) if result_parts else "No output"
+
+
+def check_for_final_answer(response: str, repl_env, logger) -> str | None:
+    """Check if response contains a final answer."""
+    # Use the new find_final_answer function which handles both FINAL and FINAL_VAR
+    return find_final_answer(response, environment=repl_env)
+
+
+def convert_context_for_repl(context):
+    """
+    Convert REPL context to either some
+    """
+    if isinstance(context, dict):
+        context_data = context
+        context_str = None
+    elif isinstance(context, str):
+        context_data = None
+        context_str = context
+    elif isinstance(context, list):
+        if len(context) > 0 and isinstance(context[0], dict):
+            if "content" in context[0]:
+                context_data = [msg.get("content", "") for msg in context]
+            else:
+                context_data = context
+            context_str = None
+        else:
+            context_data = context
+            context_str = None
+    else:
+        context_data = context
+        context_str = None
+
+    return context_data, context_str