Update app.py

app.py

@@ -5,21 +5,23 @@ import requests
 import torch
 import uvicorn
 import nest_asyncio
-import os
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from sentence_transformers import SentenceTransformer, models
 import gradio as gr
+import os
 
 ############################################
 # Configuration
 ############################################
 
+
 HF_TOKEN = os.environ.get("HF_TOKEN")
 GITHUB_TOKEN = os.environ.get("GITHUB_TOKEN")
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
 
+
 ############################################
 # GitHub API Functions
 ############################################
@@ -67,8 +69,7 @@ def preprocess_text(text: str) -> str:
 
 def load_embedding_model(model_name: str = 'huggingface/CodeBERTa-small-v1') -> SentenceTransformer:
     transformer_model = models.Transformer(model_name)
-    pooling_model = models.Pooling(transformer_model.get_word_embedding_dimension(),
-                                   pooling_mode_mean_tokens=True)
+    pooling_model = models.Pooling(transformer_model.get_word_embedding_dimension(), pooling_mode_mean_tokens=True)
     model = SentenceTransformer(modules=[transformer_model, pooling_model])
     return model
 
@@ -92,6 +93,7 @@ def generate_prompt(query: str, context_snippets: list) -> str:
         instruction = "Provide an extremely detailed and thorough explanation of at least 500 words."
     else:
         instruction = "Answer concisely."
+    
     prompt = (
         f"Below is some context from a GitHub repository:\n\n"
         f"{context}\n\n"
@@ -100,23 +102,129 @@ def generate_prompt(query: str, context_snippets: list) -> str:
     )
     return prompt
 
-# Gemini Flash integration using the new API.
+
 def get_gemini_flash_response(prompt: str) -> str:
-    import google.generativeai as palm
-    from google.generativeai import types
-    palm.configure(api_key=GEMINI_API_KEY)
-    response = palm.generate_text(
+    from google import genai 
+    from google.genai import types
+    client = genai.Client(api_key=GEMINI_API_KEY)
+
+
+    
+    response = client.models.generate_content(
         model="gemini-2.0-flash",
         contents=[prompt],
-        temperature=0.1,
-        max_output_tokens=500
+        config=types.GenerateContentConfig(
+        max_output_tokens=500,
+        temperature=0.1
+        )
     )
+    
     return response.text
 
+
+
+
+
+# ############################################
+# # Gradio Interface Functions
+# ############################################
+
+# # For file content retrieval, we now use the file path directly.
+# def get_file_content_for_choice(github_url: str, file_path: str):
+#     try:
+#         owner, repo = extract_repo_info(github_url)
+#     except Exception as e:
+#         return str(e)
+#     content = get_file_content(owner, repo, file_path)
+#     return content, file_path
+
+# def chat_with_file(github_url: str, file_path: str, user_query: str):
+#     # Retrieve file content using the file path directly.
+#     result = get_file_content_for_choice(github_url, file_path)
+#     if isinstance(result, str):
+#         return result  # Return error message if occurred.
+#     file_content, selected_file = result
+    
+#     # Preprocess file content and extract context.
+#     preprocessed = preprocess_text(file_content)
+#     context_snippet = preprocessed[:5000]  # Use first 1000 characters as context.
+    
+#     # Generate the prompt based on context and user query.
+#     prompt = generate_prompt(user_query, [context_snippet])
+    
+#     # Use Gemini Flash to generate a response.
+#     llm_response = get_gemini_flash_response(prompt)
+    
+#     return f"File: {selected_file}\n\nLLM Response:\n{llm_response}"
+
+
+# def load_repo_contents_backend(github_url: str):
+#     try:
+#         owner, repo = extract_repo_info(github_url)
+#     except Exception as e:
+#         return f"Error: {str(e)}"
+#     repo_data = get_repo_metadata(owner, repo)
+#     default_branch = repo_data.get("default_branch", "main")
+#     tree_data = get_repo_tree(owner, repo, default_branch)
+#     if "tree" not in tree_data:
+#         return "Error: Could not fetch repository tree."
+#     file_list = [item["path"] for item in tree_data["tree"] if item["type"] == "blob"]
+#     return file_list
+
 ############################################
-# Gradio Interface Functions
+# Gradio Interface Setup
 ############################################
 
+# with gr.Blocks() as demo:
+#     gr.Markdown("# RepoChat - Chat with Repository Files")
+    
+#     with gr.Row():
+#         with gr.Column(scale=1):
+#             gr.Markdown("### Repository Information")
+#             github_url_input = gr.Textbox(label="GitHub Repository URL", placeholder="https://github.com/username/repository")
+#             load_repo_btn = gr.Button("Load Repository Contents")
+#             # Dropdown with choices as file paths; default value is empty.
+#             file_dropdown = gr.Dropdown(label="Select a File", interactive=True, value="", choices=[])
+#             # repo_content_output = gr.Textbox(label="File Content", interactive=False, lines=30)
+#             repo_content_output = gr.Chatbot(label="Chat Conversation")
+#         with gr.Column(scale=2):
+#             gr.Markdown("### Chat Interface")
+#             chat_query_input = gr.Textbox(label="Your Query", placeholder="Type your query here")
+#             chat_output = gr.Textbox(label="Chatbot Response", interactive=False, lines=10)
+#             chat_btn = gr.Button("Send Query")
+    
+#     # Callback: Update file dropdown choices.
+#     def update_file_dropdown(github_url):
+#         files = load_repo_contents_backend(github_url)
+#         if isinstance(files, str):  # Error message
+#             print("Error loading files:", files)
+#             return gr.update(choices=[], value="")
+#         print("Files loaded:", files)
+#         # Do not pre-select any file (empty value)
+#         return gr.update(choices=files, value="")
+    
+#     load_repo_btn.click(fn=update_file_dropdown, inputs=[github_url_input], outputs=[file_dropdown])
+    
+#     # Callback: Update repository content when a file is selected.
+#     def update_repo_content(github_url, file_choice):
+#         if not file_choice:
+#             return "No file selected."
+#         content, _ = get_file_content_for_choice(github_url, file_choice)
+#         return content
+    
+#     file_dropdown.change(fn=update_repo_content, inputs=[github_url_input, file_dropdown], outputs=[repo_content_output])
+    
+#     # Callback: Process chat query.
+#     def process_chat(github_url, file_choice, chat_query):
+#         if not file_choice:
+#             return "Please select a file first."
+#         return chat_with_file(github_url, file_choice, chat_query)
+    
+#     chat_btn.click(fn=process_chat, inputs=[github_url_input, file_dropdown, chat_query_input], outputs=[chat_output])
+    
+# demo.launch(share=True)
+
+
 def get_file_content_for_choice(github_url: str, file_path: str):
     try:
         owner, repo = extract_repo_info(github_url)
@@ -131,7 +239,7 @@ def chat_with_file(github_url: str, file_path: str, user_query: str):
         return result  # Return error message if occurred.
     file_content, selected_file = result
     preprocessed = preprocess_text(file_content)
-    context_snippet = preprocessed[:1000]  # Use the first 1000 characters as context.
+    context_snippet = preprocessed[:1000]  # Use the first 1000 characters as context
     prompt = generate_prompt(user_query, [context_snippet])
     llm_response = get_gemini_flash_response(prompt)
     return f"File: {selected_file}\n\nLLM Response:\n{llm_response}"
@@ -162,7 +270,9 @@ with gr.Blocks() as demo:
             github_url_input = gr.Textbox(label="GitHub Repository URL", placeholder="https://github.com/username/repository")
             load_repo_btn = gr.Button("Load Repository Contents")
             file_dropdown = gr.Dropdown(label="Select a File", interactive=True, choices=[])
-            repo_content_output = gr.Textbox(label="File Content", interactive=False, lines=20)
+            # repo_content_output = gr.Textbox(label="File Content", interactive=False, lines=20)
+            repo_content_output = gr.Chatbot(label="File Content")
+        
         with gr.Column(scale=2):
             gr.Markdown("### Chat Interface")
             chat_output = gr.Chatbot(label="Chat Conversation")
@@ -174,17 +284,23 @@ with gr.Blocks() as demo:
 
     def update_file_dropdown(github_url):
         files = load_repo_contents_backend(github_url)
-        if isinstance(files, str):
-            print("Error loading files:", files)
-            return gr.update(choices=[], value=None)
-        print("Files loaded:", files)
-        return gr.update(choices=files, value=None)
+        if isinstance(files, str):  # Error message
+            return gr.update(choices=[]), files
+        return gr.update(choices=files), ""
+
+    # def update_repo_content(github_url, file_choice):
+    #     if not file_choice:
+    #         return "No file selected."
+    #     content, _ = get_file_content_for_choice(github_url, file_choice)
+    #     return content
 
     def update_repo_content(github_url, file_choice):
         if not file_choice:
-            return "No file selected."
+            return [("System", "No file selected.")]
         content, _ = get_file_content_for_choice(github_url, file_choice)
-        return content
+    # Wrap the content in a tuple so it displays like a chat message.
+        return [("File Content", content)]
+
 
     def process_chat(github_url, file_choice, chat_query, history):
         if not file_choice:
@@ -194,7 +310,7 @@ with gr.Blocks() as demo:
         history.append((chat_query, response))
         return history, history
 
-    load_repo_btn.click(fn=update_file_dropdown, inputs=[github_url_input], outputs=[file_dropdown])
+    load_repo_btn.click(fn=update_file_dropdown, inputs=[github_url_input], outputs=[file_dropdown, repo_content_output])
     file_dropdown.change(fn=update_repo_content, inputs=[github_url_input, file_dropdown], outputs=[repo_content_output])
     chat_btn.click(fn=process_chat, inputs=[github_url_input, file_dropdown, chat_query_input, conversation_history], outputs=[chat_output, conversation_history])