Update app.py

app.py

@@ -19,6 +19,8 @@ import os
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 GITHUB_TOKEN = os.environ.get("GITHUB_TOKEN")
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+
 
 
 ############################################
@@ -134,39 +136,62 @@ def generate_prompt(query: str, context_snippets: list) -> str:
     
 #     return answer
 
-def get_llm_response(prompt: str, model_name: str = "EleutherAI/gpt-neo-125M", max_new_tokens: int = None) -> str:
-    if max_new_tokens is None:
-        max_new_tokens = 256  # You can adjust this value as needed.
+# def get_llm_response(prompt: str, model_name: str = "EleutherAI/gpt-neo-125M", max_new_tokens: int = None) -> str:
+#     if max_new_tokens is None:
+#         max_new_tokens = 256  # You can adjust this value as needed.
     
-    torch.cuda.empty_cache()
+#     torch.cuda.empty_cache()
     
-    # Load the tokenizer and model for GPT-Neo 125M.
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        device_map="auto",
-        use_safetensors=False,
-        torch_dtype=torch.float32  # Using default precision since model is small.
-    )
+#     # Load the tokenizer and model for GPT-Neo 125M.
+#     tokenizer = AutoTokenizer.from_pretrained(model_name)
+#     model = AutoModelForCausalLM.from_pretrained(
+#         model_name,
+#         device_map="auto",
+#         use_safetensors=False,
+#         torch_dtype=torch.float32  # Using default precision since model is small.
+#     )
+    
+#     text_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)
+#     outputs = text_gen(
+#         prompt, 
+#         max_new_tokens=max_new_tokens, 
+#         do_sample=True, 
+#         temperature=0.9,       # Increased temperature
+#         top_p=0.9,             # Using nucleus sampling
+#         top_k=50               # Limit to top 50 tokens per step
+#     )
+#     full_response = outputs[0]['generated_text']
     
-    text_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    outputs = text_gen(
-        prompt, 
-        max_new_tokens=max_new_tokens, 
-        do_sample=True, 
-        temperature=0.9,       # Increased temperature
-        top_p=0.9,             # Using nucleus sampling
-        top_k=50               # Limit to top 50 tokens per step
+#     marker = "Answer:"
+#     if marker in full_response:
+#         answer = full_response.split(marker, 1)[1].strip()
+#     else:
+#         answer = full_response.strip()
+    
+#     return answer
+
+def get_gemini_flash_response(prompt: str) -> str:
+    from google import genai
+    from google.genai import types
+
+    # Create a Gemini client using the API key from the environment.
+    client = genai.Client(api_key=GEMINI_API_KEY)
+    
+    # Configure generation settings.
+    config = types.GenerateContentConfig(
+        max_output_tokens=500,  # Adjust as needed.
+        temperature=0.1         # Lower temperature for more deterministic responses.
     )
-    full_response = outputs[0]['generated_text']
     
-    marker = "Answer:"
-    if marker in full_response:
-        answer = full_response.split(marker, 1)[1].strip()
-    else:
-        answer = full_response.strip()
+    # Send the prompt to the Gemini-2.0-flash model.
+    response = client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents=[prompt],
+        config=config
+    )
     
-    return answer
+    return response.text
+
 
 
 ############################################
@@ -183,16 +208,25 @@ def get_file_content_for_choice(github_url: str, file_path: str):
     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  # Error message
+        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[:1000]  # use first 1000 characters as context
+    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])
-    llm_response = get_llm_response(prompt)
+    
+    # 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)