Upload 2 files

app.py

@@ -1,64 +1,900 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+
+#!/usr/bin/env python3
+# import gradio as gr
+
+import json
+import logging
+import os
+import traceback
+from pathlib import Path
+from urllib.parse import urlparse
+from typing import Dict, Any, List, Set
+from git import Repo
+import io
+
+import torch
+import numpy as np
+import faiss
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from sentence_transformers import SentenceTransformer, util
+from huggingface_hub import snapshot_download
+import os
+from openai import AzureOpenAI
+import requests
+import re
+
+import matplotlib.pyplot as plt
+from sklearn.manifold import TSNE
+from sklearn.cluster import KMeans
+import plotly.graph_objects as go
+import plotly.express as px
+import random
+
+from sklearn.cluster import AgglomerativeClustering
+
+
+def load_env():
+    from dotenv import load_dotenv
+    env_path = Path(__file__).parent.parent / '.env'
+    load_dotenv(dotenv_path=env_path)
+
+load_env()
+
+# Centralized env parameters
+HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
+GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
+AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
+AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
+MODEL_NAME = "gpt-4o-mini"
+DEPLOYMENT = "gpt-4o-mini"
+API_VERSION = "2024-12-01-preview"
+
+FILE_REGEX = re.compile(r"^diff --git a/(.+?) b/(.+)")
+LINE_HUNK = re.compile(r"@@ -(?P<old_start>\d+),(?P<old_len>\d+) \+(?P<new_start>\d+),(?P<new_len>\d+) @@")
+
+# Configure logging to capture all output
+log_stream = io.StringIO()
+log_handler = logging.StreamHandler(log_stream)
+log_handler.setLevel(logging.INFO)
+log_formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
+log_handler.setFormatter(log_formatter)
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s %(message)s",
+    handlers=[log_handler, logging.StreamHandler()]
 )
+logger = logging.getLogger(__name__)
+
+class InferenceContext:
+    def __init__(self, repo_url: str):
+        self.repo_url = repo_url
+        owner, name = self._parse_owner_repo(repo_url)
+        self.repo_id = f"{owner}/{name}"
+        self.repo_dir = f"{owner}-{name}"
+        self.hf_repo_id = "kotlarmilos/repository-learning"
+        
+        # Local paths for downloaded models
+        self.base = Path("artifacts") / self.repo_dir
+        self.model_dirs = {
+            'fine_tune': self.base / 'fine_tune',
+            'contrastive': self.base / 'contrastive',
+            'index': self.base / 'index'
+        }
+        self.code_dir = self.base / 'code'
+
+        # Create directories
+        for d in (*self.model_dirs.values(), self.code_dir):
+            d.mkdir(parents=True, exist_ok=True)
+
+    @staticmethod
+    def _parse_owner_repo(url: str) -> tuple[str, str]:
+        parts = urlparse(url).path.strip("/").split("/")
+        if len(parts) < 2:
+            raise ValueError(f"Invalid GitHub URL: {url}")
+        return parts[-2], parts[-1]
+
+
+class InferencePipeline:
+    def __init__(self, ctx: InferenceContext):
+        self.ctx = ctx
+        self.tokenizer = None
+        self.llm = None
+        self.embedder = None
+        self.faiss_index = None
+        self.faiss_metadata = None
+
+        self.download_artifacts()
+        self.load_models()
+
+    def download_artifacts(self):
+        """Download models and index from Hugging Face if they don't exist locally."""
+
+        self.repo_files = self._clone_or_pull()
+            
+        snapshot_download(
+            repo_id=self.ctx.hf_repo_id,
+            allow_patterns=f"{self.ctx.repo_dir}/**",
+            local_dir=str(self.ctx.base.parent),
+            local_dir_use_symlinks=False,
+            token=HUGGINGFACE_HUB_TOKEN
+        )
+            
+        logger.info("All artifacts download complete.")
+
+    def _clone_or_pull(self) -> bool:
+        dest = self.ctx.code_dir
+        git_dir = dest / ".git"
+        if git_dir.exists():
+            Repo(dest).remotes.origin.pull()
+            logger.info("Pulled latest code into %s", dest)
+        else:
+            Repo.clone_from(self.ctx.repo_url, dest)
+            logger.info("Cloned repo %s into %s", self.ctx.repo_url, dest)
+        
+        return [str(f.relative_to(dest)) for f in dest.rglob("*") if f.is_file()]
+
+    def load_models(self):
+        """Load the fine-tuned LLM model."""
+        self.tokenizer = AutoTokenizer.from_pretrained(self.ctx.model_dirs['fine_tune'])
+        self.local_llm = AutoModelForCausalLM.from_pretrained(
+            self.ctx.model_dirs['fine_tune'], 
+            device_map="auto", 
+            torch_dtype=torch.bfloat16
+        )
+
+        self.enterprise_llm = AzureOpenAI(
+            api_version=API_VERSION,
+            azure_endpoint=AZURE_OPENAI_ENDPOINT,
+            api_key=AZURE_OPENAI_API_KEY,
+        )
+
+        self.embedder = SentenceTransformer(str(self.ctx.model_dirs['contrastive']))
+
+        self.faiss_index = faiss.read_index(str(self.ctx.model_dirs['index'] / "index.faiss"))
+        self.faiss_metadata = json.loads((self.ctx.model_dirs['index'] / "metadata.json").read_text())
+        logger.info("FAISS index loaded successfully")
+
+    def _extract_pr_data(self, pr_url: str) -> dict:
+        """
+        Collect PR data using GitHub API.
+        """
+
+        match = re.search(r'/pull/(\d+)', pr_url)
+        pr_number = int(match.group(1))
+        
+        pr_url = f"https://api.github.com/repos/{self.ctx.repo_id}/pulls/{pr_number}"
+        comments_url = f"https://api.github.com/repos/{self.ctx.repo_id}/pulls/{pr_number}/comments"
+
+        headers = {}
+        headers["Authorization"] = f"token {GITHUB_TOKEN}"
+        headers["Accept"] = "application/vnd.github.v3+json"
+    
+        try:
+            logger.info(f"Fetching PR #{pr_number} details...")
+            pr_response = requests.get(pr_url, headers=headers)
+            pr_response.raise_for_status()
+            pr_data = pr_response.json()
+
+            logger.info(f"Fetching PR #{pr_number} review comments...")
+            comments_response = requests.get(comments_url, headers=headers)
+            comments_response.raise_for_status()
+            comments_data = comments_response.json()
+            
+            grouped = {}
+            for comment in comments_data:
+                hunk = comment.get("diff_hunk", "")
+                grouped.setdefault(hunk, []).append(comment.get("body", ""))
+
+            review_comments = [
+                {"diff_hunk": hunk, "comments": comments}
+                for hunk, comments in grouped.items()
+            ]
+            
+            logger.info(f"Fetching PR #{pr_number} diff...")
+            diff_headers = headers.copy()
+            diff_headers["Accept"] = "application/vnd.github.v3.diff"
+            diff_response = requests.get(pr_url, headers=diff_headers)
+            diff_response.raise_for_status()
+
+            parsed_diff = self.parse_diff_with_lines(diff_response.text)
+
+            result = {
+                "title": pr_data.get("title", ""),
+                "body": pr_data.get("body", ""),
+                "review_comments": review_comments,
+                "diff": diff_response.text,
+                "changed_files": list(parsed_diff['changed_files']),
+                "diff_hunks": parsed_diff['diff_hunks']
+            }
+
+            logger.info(f"Successfully collected PR #{pr_number} data")
+            return result
+            
+        except Exception as e:
+            logger.error(f"Error processing PR #{pr_number} data: {e}")
+            raise
+
+    def parse_diff_with_lines(self, diff_text: str) -> Dict[str, Any]:
+        lines = diff_text.splitlines()
+        
+        result = {
+            'changed_files': set(),
+            'diff_hunks': {}
+        }
+        
+        current_file = None
+        current_hunk_content = []
+        current_line_range = None
+        file_header_lines = []
+        
+        for line in lines:
+            # Check if this is a new file header
+            file_match = FILE_REGEX.match(line)
+            if file_match:
+                # Save previous file data if exists
+                if current_file and current_hunk_content and current_line_range:
+                    if current_file not in result['diff_hunks']:
+                        result['diff_hunks'][current_file] = []
+                    result['diff_hunks'][current_file].append({
+                        'line_range': current_line_range,
+                        'content': '\n'.join(current_hunk_content)
+                    })
+                
+                # Start new file
+                current_file = file_match.group(2)  # Use the 'b/' file path (new file)
+                result['changed_files'].add(current_file)
+                file_header_lines = [line]
+                current_hunk_content = []
+                current_line_range = None
+                
+            elif current_file:  # Only process if we're inside a file
+                # Check for hunk headers to extract line ranges
+                hunk_match = LINE_HUNK.match(line)
+                if hunk_match:
+                    # Save previous hunk if exists
+                    if current_hunk_content and current_line_range:
+                        if current_file not in result['diff_hunks']:
+                            result['diff_hunks'][current_file] = []
+                        result['diff_hunks'][current_file].append({
+                            'line_range': current_line_range,
+                            'content': '\n'.join(current_hunk_content)
+                        })
+                    
+                    # Start new hunk
+                    old_start = int(hunk_match.group('old_start'))
+                    old_len = int(hunk_match.group('old_len'))
+                    new_start = int(hunk_match.group('new_start'))
+                    new_len = int(hunk_match.group('new_len'))
+                    
+                    # Calculate the range of changed lines
+                    if new_len > 0:
+                        line_start = new_start
+                        line_end = new_start + new_len - 1
+                        current_line_range = (line_start, line_end)
+                    else:
+                        current_line_range = (new_start, new_start)
+                    
+                    # Start fresh hunk content with file headers and current hunk header
+                    current_hunk_content = file_header_lines + [line]
+                else:
+                    # Add content line to current hunk
+                    if current_hunk_content is not None:
+                        current_hunk_content.append(line)
+        
+        # Save the last hunk data
+        if current_file and current_hunk_content and current_line_range:
+            if current_file not in result['diff_hunks']:
+                result['diff_hunks'][current_file] = []
+            result['diff_hunks'][current_file].append({
+                'line_range': current_line_range,
+                'content': '\n'.join(current_hunk_content)
+            })
+        
+        return result
+    
+
+    def analyze_file_similarity(self, changed_files: List[str]) -> Dict[str, Any]:
+        result = {
+            'similar_file_groups': [],
+            'anomalous_files': [],
+            'analysis_summary': {
+                'total_files': len(changed_files),
+                'num_groups': 0,
+                'num_anomalies': 0,
+                'avg_group_size': 0
+            }
+        }
+        
+        # Handle edge cases
+        if len(changed_files) == 0:
+            logger.info("No changed files to analyze")
+            return result
+        
+        if len(changed_files) == 1:
+            logger.info(f"Only one file changed: {changed_files[0]} - no similarity analysis needed")
+            result['analysis_summary']['num_anomalies'] = 1
+            result['anomalous_files'].append({
+                'file': changed_files[0],
+                'reason': 'single_file',
+                'max_similarity_to_others': 0.0,
+                'most_similar_file': None,
+                'is_anomaly': False
+            })
+            return result
+        
+        # Encode all changed files
+        file_embeddings = self.embedder.encode(changed_files, convert_to_tensor=True)
+        similarity_matrix = util.pytorch_cos_sim(file_embeddings, file_embeddings)
+        
+        # Convert similarity matrix to distance matrix for clustering
+        distance_matrix = 1 - similarity_matrix.cpu().numpy()
+        
+        # Perform hierarchical clustering
+        clustering = AgglomerativeClustering(
+            n_clusters=None, 
+            distance_threshold=0.3,  # 1 - 0.7 = 0.3 (similarity threshold of 0.7)
+            metric='precomputed',
+            linkage='average'
+        )
+        
+        cluster_labels = clustering.fit_predict(distance_matrix)
+        
+        # Group files by cluster
+        clusters = {}
+        for i, label in enumerate(cluster_labels):
+            if label not in clusters:
+                clusters[label] = []
+            clusters[label].append((changed_files[i], i))  # Store file and its index
+        
+        # Process clusters to identify groups and anomalies
+        for cluster_id, files_with_indices in clusters.items():
+            files_in_cluster = [f[0] for f in files_with_indices]
+            
+            if len(files_in_cluster) > 1:
+                # This is a group of similar files
+                group_similarities = []
+                pairwise_similarities = []
+                
+                for i in range(len(files_with_indices)):
+                    for j in range(i+1, len(files_with_indices)):
+                        file_i_idx = files_with_indices[i][1]
+                        file_j_idx = files_with_indices[j][1]
+                        similarity = float(similarity_matrix[file_i_idx][file_j_idx])
+                        group_similarities.append(similarity)
+                        pairwise_similarities.append({
+                            'file1': files_with_indices[i][0],
+                            'file2': files_with_indices[j][0],
+                            'similarity': similarity
+                        })
+                
+                avg_similarity = sum(group_similarities) / len(group_similarities) if group_similarities else 0
+                min_similarity = min(group_similarities) if group_similarities else 0
+                max_similarity = max(group_similarities) if group_similarities else 0
+                
+                result['similar_file_groups'].append({
+                    'cluster_id': cluster_id,
+                    'files': files_in_cluster,
+                    'avg_similarity': avg_similarity,
+                    'min_similarity': min_similarity,
+                    'max_similarity': max_similarity,
+                    'pairwise_similarities': pairwise_similarities,
+                    'coherence': 'high' if min_similarity > 0.6 else 'medium' if min_similarity > 0.4 else 'low'
+                })
+            else:
+                # This is a singleton cluster - potentially anomalous
+                file = files_in_cluster[0]
+                file_idx = files_with_indices[0][1]
+                
+                # Calculate maximum similarity to any other file
+                max_similarity = 0
+                most_similar_file = None
+                similarities_to_others = []
+                
+                for other_idx, other_file in enumerate(changed_files):
+                    if other_idx != file_idx:
+                        similarity = float(similarity_matrix[file_idx][other_idx])
+                        similarities_to_others.append({
+                            'file': other_file,
+                            'similarity': similarity
+                        })
+                        if similarity > max_similarity:
+                            max_similarity = similarity
+                            most_similar_file = other_file
+                
+                result['anomalous_files'].append({
+                    'file': file,
+                    'cluster_id': cluster_id,
+                    'max_similarity_to_others': max_similarity,
+                    'most_similar_file': most_similar_file,
+                    'similarities_to_others': similarities_to_others,
+                    'is_anomaly': max_similarity < 0.5,  # Strong anomaly threshold
+                    'anomaly_strength': 'strong' if max_similarity < 0.3 else 'medium' if max_similarity < 0.5 else 'weak',
+                    'reason': 'isolated_cluster'
+                })
+        
+        # Additional anomaly detection: files that are far from the group average
+        if len(changed_files) >= 3:
+            # Calculate average embedding of all changed files
+            avg_embedding = torch.mean(file_embeddings, dim=0)
+            
+            # Find files that are far from the average
+            for i, file in enumerate(changed_files):
+                file_embedding = file_embeddings[i]
+                similarity_to_avg = float(util.pytorch_cos_sim(file_embedding.unsqueeze(0), avg_embedding.unsqueeze(0))[0][0])
+                
+                # Check if this file is already in anomalous_files
+                existing_anomaly = next((a for a in result['anomalous_files'] if a['file'] == file), None)
+                
+                if existing_anomaly:
+                    # Update existing anomaly record
+                    existing_anomaly['similarity_to_group_avg'] = similarity_to_avg
+                    existing_anomaly['is_strong_anomaly'] = (
+                        similarity_to_avg < 0.4 and existing_anomaly['max_similarity_to_others'] < 0.5
+                    )
+                    if existing_anomaly['is_strong_anomaly']:
+                        existing_anomaly['anomaly_strength'] = 'very_strong'
+                elif similarity_to_avg < 0.4:  # Low similarity to group average
+                    # Calculate similarities to all other files
+                    similarities_to_others = []
+                    max_sim = 0
+                    most_sim_file = None
+                    
+                    for j, other_file in enumerate(changed_files):
+                        if i != j:
+                            sim = float(similarity_matrix[i][j])
+                            similarities_to_others.append({
+                                'file': other_file,
+                                'similarity': sim
+                            })
+                            if sim > max_sim:
+                                max_sim = sim
+                                most_sim_file = other_file
+                    
+                    result['anomalous_files'].append({
+                        'file': file,
+                        'cluster_id': None,
+                        'max_similarity_to_others': max_sim,
+                        'most_similar_file': most_sim_file,
+                        'similarities_to_others': similarities_to_others,
+                        'similarity_to_group_avg': similarity_to_avg,
+                        'is_anomaly': True,
+                        'is_strong_anomaly': max_sim < 0.5,
+                        'anomaly_strength': 'very_strong' if max_sim < 0.3 else 'strong' if max_sim < 0.5 else 'medium',
+                        'reason': 'distant_from_group_average'
+                    })
+        
+        # Update analysis summary
+        result['analysis_summary']['num_groups'] = len(result['similar_file_groups'])
+        result['analysis_summary']['num_anomalies'] = len(result['anomalous_files'])
+        
+        if result['similar_file_groups']:
+            total_files_in_groups = sum(len(group['files']) for group in result['similar_file_groups'])
+            result['analysis_summary']['avg_group_size'] = total_files_in_groups / len(result['similar_file_groups'])
+        
+        # Log results
+        logger.info(f"File similarity analysis complete:")
+        logger.info(f"  Total files: {result['analysis_summary']['total_files']}")
+        logger.info(f"  Similar groups: {result['analysis_summary']['num_groups']}")
+        logger.info(f"  Anomalous files: {result['analysis_summary']['num_anomalies']}")
+        
+        for i, group in enumerate(result['similar_file_groups']):
+            logger.info(f"  Group {i+1} ({group['coherence']} coherence): {group['files']} (avg: {group['avg_similarity']:.3f})")
+        
+        for anomaly in result['anomalous_files']:
+            logger.info(f"  {anomaly['anomaly_strength'].upper()} ANOMALY: {anomaly['file']} (reason: {anomaly['reason']}, max_sim: {anomaly['max_similarity_to_others']:.3f})")
+        
+        return result
+
+    # TODO: Add local LLM reasoning
+    # def generate_llm_response(self, prompt: str, max_new_tokens: int = 256) -> str:
+    #     """Generate response using the fine-tuned LLM."""
+    #     if not self.tokenizer or not self.local_llm:
+    #         raise ValueError("LLM not loaded. Call load_llm() first.")
+        
+    #     inputs = self.tokenizer(prompt, return_tensors="pt").to(self.local_llm.device)
+    #     outputs = self.local_llm.generate(
+    #         **inputs, 
+    #         max_new_tokens=max_new_tokens, 
+    #         pad_token_id=self.tokenizer.eos_token_id
+    #     )
+    #     return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    def search_code_snippets(self, diff_hunks) -> list:
+        metadata_file = self.ctx.model_dirs["index"] / "metadata.json"
+        with open(metadata_file, 'r', encoding='utf-8') as f:
+            metadata = json.load(f)
+        
+        result = []
+        
+        # Process each file's diff hunks
+        for file_path, hunks in diff_hunks.items():
+            logger.info(f"Searching functions for file: {file_path}")
+            
+            for hunk in hunks:
+                line_range = hunk.get('line_range')
+                if not line_range:
+                    continue
+                    
+                start_line, end_line = line_range
+                logger.debug(f"Processing hunk at lines {start_line}-{end_line}")
+                
+                # Find functions that overlap with this line range
+                overlapping_functions = []
+                
+                for func_metadata in metadata:
+                    func_file = func_metadata.get('file', '')
+                    func_start = func_metadata.get('start_line')
+                    func_end = func_metadata.get('end_line')
+                    func_name = func_metadata.get('name', 'unknown')
+                    func_description = func_metadata.get('llm_description', '')
+                    
+                    # Check if this function is in the same file
+                    if func_file != file_path:
+                        continue
+                    
+                    # Check if function line range overlaps with diff hunk line range
+                    if func_start is not None and func_end is not None:
+                        # Check for overlap: function overlaps if it starts before diff ends 
+                        # and ends after diff starts
+                        if func_start <= end_line and func_end >= start_line:
+                            overlap_start = max(func_start, start_line)
+                            overlap_end = min(func_end, end_line)
+                            
+                            overlapping_functions.append({
+                                'function_name': func_name,
+                                'function_description': func_description,
+                                'function_start_line': func_start,
+                                'function_end_line': func_end,
+                                # 'overlap_start': overlap_start,
+                                # 'overlap_end': overlap_end,
+                                # 'overlap_lines': overlap_end - overlap_start + 1
+                            })
+                
+                # if len(overlapping_functions) > 0:
+                hunk_result = {
+                    'file_name': file_path,
+                    'diff_hunk': hunk.get('content', ''),
+                    'overlapping_functions': overlapping_functions
+                }
+                
+                result.append(hunk_result)
+        
+        total_hunks = sum(len(hunks) for hunks in diff_hunks.values())
+        total_functions = sum(len(entry['overlapping_functions']) for entry in result)
+        logger.info(f"Processed {total_hunks} diff hunks across {len(diff_hunks)} files, found {total_functions} overlapping functions")
+        
+        return result
+
+    def _select_files_around_changed(self, changed_files: List[str] = None, max_files: int = 500) -> List[str]:
+        """Select files to visualize, prioritizing changed files and semantically similar ones."""
+        
+        logger.info(f"Selecting {max_files} files around {len(changed_files)} changed files...")
+        
+        # Start with changed files
+        selected_files = set(changed_files)
+        
+        # Find files similar to changed files using embeddings
+        try:
+            # Encode changed files
+            changed_embeddings = self.embedder.encode(changed_files, convert_to_tensor=False)
+            
+            # Calculate target number of similar files to find
+            target_similar = min(max_files - len(changed_files), 200)  # Leave room for random files
+            
+            # Get a sample of repo files to compare against (for performance)
+            sample_size = min(2000, len(self.repo_files))
+            repo_sample = self.repo_files[:sample_size]
+            
+            # Remove already selected files from sample
+            repo_sample = [f for f in repo_sample if f not in selected_files]
+            
+            if len(repo_sample) > 0:
+                # Encode sample files
+                sample_embeddings = self.embedder.encode(repo_sample, convert_to_tensor=False, show_progress_bar=False)
+                
+                # Calculate similarities
+                similarities = []
+                for i, repo_file in enumerate(repo_sample):
+                    # Calculate max similarity to any changed file
+                    max_sim = 0
+                    for changed_emb in changed_embeddings:
+                        sim = np.dot(changed_emb, sample_embeddings[i]) / (
+                            np.linalg.norm(changed_emb) * np.linalg.norm(sample_embeddings[i])
+                        )
+                        max_sim = max(max_sim, sim)
+                    # Only add if not already selected (avoid duplicates)
+                    similarities.append((repo_file, max_sim))
+                
+                # Sort by similarity and take top ones, avoiding duplicates
+                added = 0
+                for file_path, sim in sorted(similarities, key=lambda x: x[1], reverse=True):
+                    if file_path not in selected_files:
+                        selected_files.add(file_path)
+                        added += 1
+                        if len(selected_files) >= max_files or added >= target_similar:
+                            break
+                logger.info(f"Added {len(similarities[:target_similar])} similar files to visualization")
+        
+        except Exception as e:
+            logger.warning(f"Could not compute file similarities: {e}")
+        
+        # Fill remaining slots with random files
+        remaining_slots = max_files - len(selected_files)
+        if remaining_slots > 0:
+            remaining_files = [f for f in self.repo_files if f not in selected_files]
+            random.shuffle(remaining_files)
+            for file_path in remaining_files[:remaining_slots]:
+                selected_files.add(file_path)
+        
+        result = list(selected_files)
+        logger.info(f"Selected {len(result)} files total: {len(changed_files)} changed, {len(result) - len(changed_files)} related/random")
+        return result
+
+    def create_repo_visualization(self, changed_files: List[str] = None, max_files: int = 500):
+        files_to_plot = self._select_files_around_changed(changed_files, max_files * len(changed_files))
+        logger.info(f"Creating visualization for {len(files_to_plot)} files...")
+        
+        if len(files_to_plot) < 2:
+            return self._create_dummy_plot(f"Only {len(files_to_plot)} files available")
+
+        embeddings = self.embedder.encode(files_to_plot, convert_to_tensor=False, show_progress_bar=False)
+        logger.info(f"Embeddings computed successfully: shape {getattr(embeddings, 'shape', None)}")
+
+        n = len(files_to_plot)
+        perplexity = min(30, max(1, n - 1))
+        tsne = TSNE(n_components=3, perplexity=perplexity, init='random', random_state=42)
+        reduced = tsne.fit_transform(embeddings)
+
+        fig = go.Figure()
+        
+        colors = []
+        sizes = []
+        hover_texts = []
+        
+        for i, file_path in enumerate(files_to_plot):
+            if changed_files and file_path in changed_files:
+                colors.append('red')
+            else:
+                # Color by file type
+                ext = os.path.splitext(file_path)[1].lower()
+                if ext in ['.py', '.js', '.ts', '.java', '.cpp', '.c', '.cs', '.rb', '.go', '.rs']:
+                    colors.append('blue')
+                elif ext in ['.md', '.txt', '.rst', '.doc']:
+                    colors.append('green')
+                elif ext in ['.json', '.yaml', '.yml', '.xml', '.toml', '.ini']:
+                    colors.append('orange')
+                elif ext in ['.html', '.css', '.scss', '.sass']:
+                    colors.append('purple')
+                else:
+                    colors.append('gray')
+            sizes.append(8)
+            hover_texts.append(f"{os.path.basename(file_path)}")
+        
+        fig.add_trace(go.Scatter3d(
+                    x=reduced[:, 0].tolist(),
+                    y=reduced[:, 1].tolist(),
+                    z=reduced[:, 2].tolist(),
+                    mode='markers+text',
+                    marker=dict(size=sizes, color=colors),
+                    text=[os.path.basename(f) for f in files_to_plot],
+                    hovertext=hover_texts,
+                    textposition='middle center',
+                    name='Repository Files'
+                ))
+        
+        title_text = 'Repository File Embeddings (3D t-SNE)'
+        if changed_files:
+            title_text += f' - {len(changed_files)} Changed Files Highlighted in Red'
+        
+        fig.update_layout(
+            title=title_text,
+            scene=dict(
+                xaxis_title='t-SNE 1',
+                yaxis_title='t-SNE 2',
+                zaxis_title='t-SNE 3',
+                camera=dict(eye=dict(x=1.5, y=1.5, z=1.5))
+            ),
+            width=800,
+            height=600,
+            margin=dict(r=20, b=10, l=10, t=60)
+        )
+
+        return fig
+
+def get_current_logs():
+    return log_stream.getvalue()
+
+# Pipeline
+
+pipeline = InferencePipeline(InferenceContext("https://github.com/dotnet/xharness"))
+
+def analyze_pr_streaming(pr_url):
+    log_stream.seek(0)
+    log_stream.truncate()
+    
+    response = {}
+    base_review = ""
+    final_review = ""
+    visualization = None
+    
+    data = pipeline._extract_pr_data(pr_url)
+    yield base_review, final_review, get_current_logs(), visualization
+    
+    visualization = pipeline.create_repo_visualization(list(data["changed_files"]), max_files=20)
+    yield "", "", get_current_logs(), visualization
+
+    similarity_analysis = pipeline.analyze_file_similarity(list(data["changed_files"]))
+    similar_file_groups = similarity_analysis['similar_file_groups']
+    anomalous_files = similarity_analysis['anomalous_files']
+    yield "", "", get_current_logs(), visualization
+
+    code_description = pipeline.search_code_snippets(data["diff_hunks"])
+
+    # Base prompt
+    base_prompt = f"""You are an expert code reviewer. Analyze the PR below and provide detailed feedback with specific line-by-line suggestions:
+        Provide a detailed code review including:
+        1. **Code Quality Issues**: Point out specific lines with problems
+        2. **Suggested Fixes**: Provide exact code suggestions with diff format
+        3. **Security Concerns**: Highlight any security vulnerabilities
+        4. **Performance Issues**: Identify potential performance problems
+        5. **Best Practices**: Suggest improvements following coding standards
+        6. **Testing Recommendations**: What tests should be added/modified
+
+        Format your suggestions like this for each issue:
+        **File: `filename.ext` Line X**
+        Problem: [description]
+        ```diff
+        - old code line
+        + suggested new code line
+        ```
+
+        TITLE: {data['title']}
+
+        DESCRIPTION: {data['body']}
+
+        CHANGED FILES: {', '.join(data['changed_files'])}
+        """
+    
+    similar_file_groups_formatted = []
+    for i, group in enumerate(similar_file_groups):
+        files_str = ", ".join(group['files'])
+        similar_file_groups_formatted.append(f"group {i}: {files_str}")
+    
+    anomalous_files_formatted = []
+    for anomaly in anomalous_files:
+        anomalous_files_formatted.append(f"anomaly: {anomaly['file']} (reason: {anomaly['reason']}, strength: {anomaly['anomaly_strength']})")
+
+    grounding_formatted = ""
+    for entry in code_description:
+        file_name = entry['file_name']
+        overlapping_functions = entry['overlapping_functions']
+        diff_hunk = entry['diff_hunk']
+        
+        if len(overlapping_functions) > 0:
+            grounding_formatted += f"In file {file_name}, the following changes were made: {diff_hunk}\n"
+            grounding_formatted += f"These changes affected the following functions:\n"
+            for func in overlapping_functions:
+                grounding_formatted += f"{func['function_name']} - {func['function_description']}\n"
+        else:
+            grounding_formatted += f"In file {file_name}, the following changes were made: {diff_hunk}\n"
+
+        grounding_formatted += "\n"
+
+    # Create formatted strings for f-string
+    similar_groups_text = "\n".join(similar_file_groups_formatted)
+    anomalous_files_text = "\n".join(anomalous_files_formatted)
+
+
+    # TODO: Add local LLM reasoning
+    # TODO: Add relevant files from the directory not included
+    comprehensive_prompt = f"""{base_prompt}
+        FILES THAT ARE SEMANTICALLY CLOSE CHANGED IN THIS PR: 
+        {similar_groups_text}
+
+        UNEXPECTED CHANGES IN FILES: 
+        {anomalous_files_text}
+
+        GROUNDING DATA: The following provides specific information about which functions are affected by each diff hunk:
+        {grounding_formatted}
+        """
+
+    base_prompt += f"""
+        DIFF: {data['diff']}
+    """
+
+    logger.info(f"Base prompt word count: {len(base_prompt.split())}")
+    # logger.info(f"Base prompt: {base_prompt}")
+
+    logger.info(f"Comprehensive prompt word count: {len(comprehensive_prompt.split())}")
+    # logger.info(f"Comprehensive prompt: {comprehensive_prompt}")
+
+    logger.info("Calling Azure OpenAI...")
+    yield "", "", get_current_logs(), visualization
+    
+    base_review_response = pipeline.enterprise_llm.chat.completions.create(
+        model=DEPLOYMENT,
+        messages=[
+            {"role": "system", "content": "You are an expert code reviewer. Provide thorough, constructive feedback."},
+            {"role": "user", "content": base_prompt}
+        ],
+        max_tokens=8192,
+        temperature=0.3
+    )
+    
+    base_review = base_review_response.choices[0].message.content
+    logger.info("Base review completed")
+    
+    final_review_response = pipeline.enterprise_llm.chat.completions.create(
+        model=DEPLOYMENT,
+        messages=[
+            {"role": "system", "content": "You are an expert code reviewer. Provide thorough, constructive feedback."},
+            {"role": "user", "content": comprehensive_prompt}
+        ],
+        max_tokens=8192,
+        temperature=0.3
+    )
+    
+    final_review = final_review_response.choices[0].message.content
+    logger.info("Final review completed")
+    
+    yield base_review, final_review, get_current_logs(), visualization
 
+with gr.Blocks(title="PR Code Review Assistant") as demo:
+    gr.Markdown("# PR Code Review Assistant")
+    gr.Markdown("Enter a GitHub PR URL to get comprehensive code review analysis with interactive repository visualization.")
+    
+    with gr.Row():
+        pr_url_input = gr.Textbox(
+            label="GitHub PR URL",
+            placeholder="https://github.com/owner/repo/pull/123",
+            value="https://github.com/dotnet/xharness/pull/1416"
+        )
+    
+    analyze_btn = gr.Button("Analyze PR", variant="primary")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            base_review_output = gr.Textbox(
+                label="Base Review",
+                lines=15,
+                max_lines=30,
+                interactive=False
+            )
+        
+        with gr.Column(scale=1):
+            final_review_output = gr.Textbox(
+                label="Comprehensive Review",
+                lines=15,
+                max_lines=30,
+                interactive=False
+            )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            visualization_output = gr.Plot(
+                label="Repository Files Visualization (3D)",
+                value=None
+            )
+        
+        with gr.Column(scale=1):
+            logs_output = gr.Textbox(
+                label="Analysis Logs",
+                lines=15,
+                max_lines=25,
+                interactive=False,
+                show_copy_button=True
+            )
+    
+    analyze_btn.click(
+        fn=analyze_pr_streaming,
+        inputs=[pr_url_input],
+        outputs=[base_review_output, final_review_output, logs_output, visualization_output],
+        show_progress=True
+    )
 
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -1 +1,54 @@
-huggingface_hub==0.25.2
+# Core ML dependencies
+torch>=2.0.0
+transformers>=4.30.0
+sentence-transformers>=2.2.2
+faiss-cpu>=1.7.4
+numpy>=1.21.0
+
+# Data processing
+pandas>=1.5.0
+datasets>=2.12.0
+accelerate>=0.20.0
+peft>=0.4.0
+bitsandbytes>=0.39.0
+huggingface_hub>=0.16.0
+
+# GitHub integration  
+PyGithub
+GitPython>=3.1.0
+requests>=2.28.0
+python-dotenv>=1.0.0
+tenacity>=8.2.0
+
+# Tree-sitter parsers
+tree-sitter>=0.20.0
+tree-sitter-python>=0.20.0
+tree-sitter-c>=0.20.0
+tree-sitter-cpp>=0.20.0
+tree-sitter-java>=0.20.0
+tree-sitter-c-sharp>=0.20.0
+tree-sitter-javascript>=0.20.0
+
+# Web interface
+flask>=2.3.0
+flask-session>=0.5.0
+
+# Utilities
+tqdm>=4.64.0
+schedule>=1.2.0
+openai
+
+# Development
+pytest>=7.0.0
+black>=23.0.0
+flake8>=6.0.0
+
+# Transformers library for NLP tasks
+transformers
+matplotlib>=3.5.0
+scikit-learn>=1.1.0
+gradio>=4.0.0
+
+# Visualization dependencies
+plotly>=5.0.0
+openai