Update buffalo_rag/model/rag.py

buffalo_rag/model/rag.py

@@ -31,11 +31,9 @@ class BuffaloRAG:
                 query: str, 
                 k: int = 5,
                 filter_categories: Optional[List[str]] = None) -> List[Dict[str, Any]]:
-        """Retrieve relevant chunks for a query."""
         return self.vector_store.hybrid_search(query, k=k, filter_categories=filter_categories)
     
     def format_context(self, results: List[Dict[str, Any]]) -> str:
-        """Concatenate retrieved passages into context."""
         ctx = []
         for i, r in enumerate(results, start=1):
             c = r["chunk"]
@@ -47,7 +45,6 @@ class BuffaloRAG:
         return "\n".join(ctx)
     
     def generate_response(self, query: str, context: str) -> str:
-        """Generate response using the language model with error handling."""
         prompt = f"""You are a friendly and professional counselor for international students at the University at Buffalo. Respond to the student's query in a supportive, detailed, and well-structured manner.
 
                 For your responses:
@@ -82,24 +79,20 @@ class BuffaloRAG:
             return completion.choices[0].message.content
         except Exception as e:
             print(f"Error during generation: {str(e)}")
-            # Fallback response
             return "I'm sorry, I encountered an issue generating a response. Please try asking your question in a different way or contact UB International Student Services directly for assistance."
     
     def answer(self, 
               query: str, 
               k: int = 5,
               filter_categories: Optional[List[str]] = None) -> Dict[str, Any]:
-        """End-to-end RAG pipeline."""
-        # Retrieve relevant chunks
+        
         results = self.retrieve(query, k=k, filter_categories=filter_categories)
         
-        # Format context
+        
         context = self.format_context(results)
         
-        # Generate response
         response = self.generate_response(query, context)
         
-        # Return response and sources
         return {
             'query': query,
             'response': response,
@@ -111,15 +104,4 @@ class BuffaloRAG:
                 }
                 for result in results
             ]
-        }
-
-# Example usage
-if __name__ == "__main__":
-    rag = BuffaloRAG(model_name="1bitLLM/bitnet_b1_58-large")
-    response = rag.answer("How do I apply for OPT?")
-    
-    print(f"Query: {response['query']}")
-    print(f"Response: {response['response']}")
-    print("\nSources:")
-    for source in response['sources']:
-        print(f"- {source['title']} (Score: {source['score']:.4f})") 
+        }