Added G2.0 changes

backend/api/main.py

@@ -33,7 +33,6 @@ templates = Jinja2Templates(directory="backend/api/templates")
 
 load_dotenv()
 
-
 logger = initialize_logger()
 
 # get current file path using Path
@@ -54,8 +53,8 @@ embedding_model = get_embedding_model(EmbeddingModel, embedding_model_config)
 
 # Create vector db model object
 vector_db_config = MilvusVectorDatabaseConfig(
-    db_path=app_config["vector_database"]["db_path"],
-    collection_name=app_config["vector_database"]["collection_name"],
+    db_path=app_config["vector_database"]["db_path"] + env_variables["MILVUS_DB"] + "_milvus.db",
+    collection_name=env_variables["MILVUS_DB"],
     vector_dimensions=app_config["vector_database"]["dimensions"],
     drop_if_exists=False,
 )
@@ -63,25 +62,28 @@ vector_db = create_vector_database(MilvusVectorDatabase, vector_db_config)
 
 # Create generative model object
 gemini_generative_model_config = GeminiModelConfig(
-    model_name=app_config["gemini_generative_model"]["model_name"],
+    model_name=env_variables["GOOGLE_GEMINI_MODEL"],
     api_keys=[env_variables["GOOGLE_GEMINI_API_KEY"], env_variables["GOOGLE_GEMINI_BACKUP_API_KEY"]],
-    temperature=app_config["gemini_generative_model"]["temperature"],
+    temperature=int(env_variables["MODEL_TEMPERATURE"]),
 )
 gemini_generative_model = get_generative_model(GeminiModel, gemini_generative_model_config)
 
 openai_generative_model_config = OpenAIModelConfig(
-    model_name=app_config["openai_generative_model"]["model_name"],
+    model_name=env_variables["OPENAI_MODEL"],
     api_key=env_variables["OPENAI_API_KEY"],
-    temperature=app_config["openai_generative_model"]["temperature"],
+    temperature=int(env_variables["MODEL_TEMPERATURE"]),
 )
 openai_generative_model = get_generative_model(OpenAIModel, openai_generative_model_config)
 
+default_project_name = env_variables["GALILEO_PROJECT_NAME"]
+default_logstream_name = env_variables["GALILEO_LOGSTREAM_NAME"]
+default_protect_stage_name = env_variables["GALILEO_PROTECT_STAGE_NAME"]
+default_dataset_name = env_variables["GALILEO_DATASET_NAME"]
+
 # Create Galileo platform object
 galileo_platform_config = GalileoPlatformConfig(
-    evaluate_project_name=app_config["galileo_platform"]["evaluate_project_name"],
-    observe_project_name=app_config["galileo_platform"]["observe_project_name"],
-    protect_project_name=app_config["galileo_platform"]["protect_project_name"],
-    protect_stage_name=app_config["galileo_platform"]["protect_stage_name"],
+    protect_project_name=env_variables["GALILEO_PROJECT_NAME"],
+    protect_stage_name=default_protect_stage_name,
 )
 galileo_platform = GalileoPlatform(galileo_platform_config)
 
@@ -89,8 +91,8 @@ galileo_platform = GalileoPlatform(galileo_platform_config)
 rag_application_config = RAGApplicationConfig(
     embedding_model=embedding_model,
     vector_db=vector_db,
-    # gemini_generative_model=gemini_generative_model,
-    generative_model=openai_generative_model,
+    generative_model=gemini_generative_model,
+    # generative_model=openai_generative_model,
     galileo_platform=galileo_platform,
 )
 rag_app = RAGApplication(rag_application_config)
@@ -98,27 +100,36 @@ rag_app = RAGApplication(rag_application_config)
 
 @app.get("/", response_class=HTMLResponse)
 async def read_root(request: Request):
-    return templates.TemplateResponse("index.html", {"request": request})
-
-# TODO: Nikhil
-# @app.post("/other-metrics")
-# async def search(
-
+    # Get default project name from environment variables
+    return templates.TemplateResponse("index.html", {
+        "request": request, 
+        "default_project_name": default_project_name,
+        "default_logstream_name": default_logstream_name,
+        "default_dataset_name": default_dataset_name
+    })
 
 @app.post("/search")
 async def search(
     query: str = Form(...),
     top_k: int = Form(5),
+    add_to_dataset: bool = Form(False),
     protection: bool = Form(False),
     hallucination_detection: bool = Form(False),
     induce_hallucination: bool = Form(False),
+    project_name: str = default_project_name,
+    logstream_name: str = default_logstream_name,
+    dataset_name: str = default_dataset_name,
 ):
+
     response, redacted_response, original_response, context_adherence_score, pii_flag = rag_app.run(
         query,
         protect_enabled=protection,
         top_k=top_k,
         hallucination_detection=hallucination_detection,
         induce_hallucination=induce_hallucination,
+        project_name=project_name,
+        logstream_name=logstream_name,
+        dataset_name=dataset_name if add_to_dataset else None,
     )
 
     # Simulate processing

backend/classes/galileo_platform.py

@@ -1,116 +1,102 @@
-from galileo_observe import ObserveWorkflows
-import galileo_protect as gp
-from pydantic import BaseModel
-import promptquality as pq
-from promptquality import CustomizedScorerName, Models
-from dotenv import load_dotenv
-import os
-from datetime import datetime
 from typing import Optional
+
+from dotenv import load_dotenv
+from pydantic import BaseModel
+
+from galileo import GalileoLogger, GalileoScorers, StageType
+from galileo.protect import invoke_protect
+from galileo.stages import create_protect_stage, get_protect_stage
+from galileo_core.schemas.protect.action import OverrideAction
+from galileo_core.schemas.protect.payload import Payload
+from galileo_core.schemas.protect.rule import Rule, RuleOperator
+from galileo_core.schemas.protect.ruleset import Ruleset
+
 load_dotenv()
 
 class GalileoPlatformConfig(BaseModel):
     """Base configuration for Galileo platform."""
-    evaluate_project_name: str
-    observe_project_name: str
     protect_project_name: str
     protect_stage_name: str
 
-
 class GalileoPlatform:
     """Implementation of Galileo Features"""
 
     def __init__(self, config: GalileoPlatformConfig):
         self.config = config
-        pq.login(api_key=os.getenv("GALILEO_API_KEY"))
-        self.evaluate_run = self.create_evaluate_run()
-        self.observe_logger = ObserveWorkflows(project_name=config.observe_project_name)
-        self.protect_stage_id = self.get_protect_stage()
+        self.protect_stage_id = self.get_protect_stage_id()
 
-    def create_evaluate_run(self):
-        """Create a Galileo Evaluate run."""
-        scorers = [
-            pq.Scorers.context_adherence_luna,
-            pq.Scorers.chunk_attribution_utilization_luna,
-            pq.Scorers.completeness_luna
-        ]
-        evaluate_run = pq.EvaluateRun(
-            project_name=self.config.evaluate_project_name,
-            scorers=scorers,
+    def get_logger(self, project_name: str, logstream_name: str):
+        """Get or create a Galileo Logger."""
+        return GalileoLogger(
+            project=project_name,
+            log_stream=logstream_name,
         )
-        return evaluate_run
 
-    def get_protect_stage(self):
+    def get_protect_stage_id(self):
         """Get or create a Galileo Protect stage."""
         try:
-            protect_project = gp.get_project(
-                project_name=self.config.protect_project_name
+            protect_stage = get_protect_stage(
+                project_name=self.config.protect_project_name,
+                stage_name=self.config.protect_stage_name,
             )
+            return protect_stage.id
         except Exception as _:
-            protect_project = gp.create_project(name=self.config.protect_project_name)
-
-        protect_project_id = protect_project.id
-
-        try:
-            protect_stage = gp.get_stage(
-                project_id=protect_project_id, stage_name=self.config.protect_stage_name
-            )
-        except Exception as _:
-            protect_stage = gp.create_stage(
-                project_id=protect_project_id,
+            protect_stage = create_protect_stage(
+                project_name=self.config.protect_project_name,
                 name=self.config.protect_stage_name,
+                stage_type=StageType.local,
+                description="Deutsche Bank RFP RAG Protect Stage"
             )
+            return protect_stage.id
 
-        return protect_stage.id
-
-    def run_protect(self, prompt: str, output: str, workflow: Optional[ObserveWorkflows] = None) -> dict:
+    def run_protect(self, input: str, output: str, logger: Optional[GalileoLogger] = None) -> dict:
         """Run Galileo Protect on input and output."""
-        response = gp.invoke(
-            payload=gp.Payload(input=prompt, output=output),
+        response = invoke_protect(
+            payload=Payload(input=input, output=output),
             prioritized_rulesets=[
-                gp.Ruleset(
+                Ruleset(
                     rules=[
-                        gp.Rule(
-                            metric=gp.RuleMetrics.context_adherence_luna,
-                            operator=gp.RuleOperator.lte,
+                        Rule(
+                            metric=GalileoScorers.context_adherence_luna,
+                            operator=RuleOperator.lte,
                             target_value=0.01,
                         ),
                     ],
-                    action=gp.OverrideAction(
+                    action=OverrideAction(
                         choices=["Sorry, the input is hallucinatory."]
                     ),
                 ),
-                gp.Ruleset(
+                Ruleset(
                     rules=[
-                        gp.Rule(
-                            metric=gp.RuleMetrics.pii,
-                            operator=gp.RuleOperator.any,
+                        Rule(
+                            metric=GalileoScorers.input_pii,
+                            operator=RuleOperator.any,
                             target_value=["email", "phone_number", "name"],
                         )
                     ],
-                    action=gp.OverrideAction(
+                    action=OverrideAction(
                         choices=["Sorry, the output contains PII."]
                     ),
                 ),
-                # gp.Ruleset(
-                #     rules=[
-                #         gp.Rule(
-                #             metric="deutsche_bank_company_pii_0",
-                #             operator=gp.RuleOperator.gte,
-                #             target_value=0.1,
-                #         )
-                #     ],
-                #     action=gp.OverrideAction(
-                #         choices=["Sorry, the output contains PII."]
-                #     ),
-                # )
+                Ruleset(
+                    rules=[
+                        Rule(
+                            metric="deutsche_bank_company_pii_scorer_0",
+                            operator=RuleOperator.gte,
+                            target_value=0.1,
+                        )
+                    ],
+                    action=OverrideAction(
+                        choices=["Sorry, the output contains PII."]
+                    ),
+                )
             ],
             stage_id=self.protect_stage_id,
         )
 
-        if workflow:
-            workflow.add_protect(
-                payload=gp.Payload(input=prompt, output=output),
+        if logger:
+            logger.add_protect_span(
+                payload=Payload(input=input, output=output),
                 response=response,
             )
 

backend/classes/rag_application.py

@@ -1,13 +1,13 @@
-from pydantic import BaseModel
-import json
 import time
-import re
-from promptquality import Models
+from typing import List, Union
+
+from galileo.datasets import create_dataset, get_dataset
+from pydantic import BaseModel
+
 from backend.classes.embedding_model import EmbeddingModel
-from backend.classes.vector_database.milvus_vector_database import MilvusVectorDatabase
 from backend.classes.galileo_platform import GalileoPlatform
 from backend.classes.generative_model import GeminiModel, OpenAIModel
-from typing import Union
+from backend.classes.vector_database.milvus_vector_database import MilvusVectorDatabase
 
 def strike(text):
     return ''.join([char + '\u0336' for char in text])
@@ -43,6 +43,7 @@ The following are the categories that need to be redacted:
 - Phone numbers
 - Email addresses
 - Names
+- Company names (Fairfield or Fairfield CDC or other variations)
 For every PII that needs to be redacted, wrap it in <pii></pii> tags.
 
 Categories: {pii_flag}
@@ -50,7 +51,7 @@ Response: {response}
 Modified Response: """
 
 
-    hallucinatory_chunks: list[str] = [
+    hallucinatory_chunks: List[str] = [
         "Fairfield CDC is issuing this RFP to select a banking partner for its ambitious new program to fund the city's first dragon-powered public transportation system.",
         "Merchant services must include psychic energy transfer gateways for multi-reality donation collection.",
         "Technological capabilities must include temporal online banking for pre-cognitive transaction approvals.",
@@ -75,14 +76,15 @@ class RAGApplication:
         top_k: int = 5,
         hallucination_detection: bool = False,
         induce_hallucination: bool = False,
+        project_name: str = None,
+        logstream_name: str = None,
+        dataset_name: str = None,
     ) -> str:
-        # Create a workflow to track this query
-        observe_workflow = self.config.galileo_platform.observe_logger.add_workflow(
-            name="RAG Workflow", input={"query": query}
-        )
+        
+        galileo_logger = self.config.galileo_platform.get_logger(project_name, logstream_name)
 
-        evaluate_workflow = self.config.galileo_platform.evaluate_run.add_workflow(
-            name="RAG Workflow", input={"query": query}
+        _ = galileo_logger.start_trace(
+            name="RAG Workflow", input=query
         )
 
         context_adherence_score = 1
@@ -97,10 +99,8 @@ class RAGApplication:
         try:
             start_time = time.time()
 
-            # Get query embedding
             query_embedding = self.config.embedding_model.encode([query])
 
-            # Get top-k similar texts
             retrieved_documents = [
                 str(text["text"])
                 for text in self.config.vector_db.search_similar_texts(
@@ -108,42 +108,27 @@ class RAGApplication:
                 )
             ]
 
-            # Log retriever step to Galileo Observe
-            observe_workflow.add_retriever(
+            galileo_logger.add_retriever_span(
                 name="Milvus Retrieval",
                 input=query,
-                documents=retrieved_documents,
+                output=retrieved_documents,
                 duration_ns=int((time.time() - start_time) * 1e9),
             )
-
-            evaluate_workflow.add_retriever(
-                name="Milvus Retrieval",
-                input=query,
-                documents=retrieved_documents,
-                # documents=[
-                # Document(content=doc, metadata={"length": len(doc)}) for doc in retrieved_documents],
-                duration_ns=int((time.time() - start_time) * 1e9),
-            )
-
             start_time = time.time()
 
             if not retrieved_documents:
                 return "There is nothing to return", redacted_result, context_adherence_score, pii_flag
             
-            # Create context by combining the retrieved documents
             context = "\n\n".join(retrieved_documents)
 
-            # Set prompt template
             prompt = (
                 self.config.prompt_template
                 if not prompt_template
                 else prompt_template
             )
 
-            # Construct prompt
             formatted_prompt = f"{prompt}\n\nQUESTION: {query}\n\nCONTEXT: {context}"
 
-            # Generate response
             result = self.config.generative_model.generate_response(
                 formatted_prompt
             )
@@ -156,59 +141,63 @@ class RAGApplication:
                     temperature=1.0,
                 )
 
-            # Log LLM call to Galileo Observe
-            observe_workflow.add_llm(
+            input_data = {
+                "question": query,
+                "context": context,
+            }
+
+            galileo_logger.add_llm_span(
                 name="Answer Generation",
-                input=retrieved_documents,
+                input=input_data,
                 output=result,
                 model=self.config.generative_model.config.model_name,
                 duration_ns=int((time.time() - start_time) * 1e9),
+                metadata={
+                    "question": query,
+                    "context": context,
+                }
             )
 
-            evaluate_workflow.add_llm(
-                # input=Message(content=prompt, role=MessageRole.user),
-                # output=Message(content=result, role=MessageRole.assistant),
-                name="Answer Generation",
-                input=prompt,
-                output=result,
-                model=Models.gpt_4o,
-                duration_ns=int((time.time() - start_time) * 1e9),
-            )
+            try:
+                row = {
+                    "input": input_data,
+                    "output": result,
+                }
+                if dataset_name:
+                    dataset = get_dataset(name=dataset_name)
+                    dataset.add_rows([row])
+            except Exception as e:
+                print(e)
+                dataset = create_dataset(name=dataset_name, content=[row])
 
             start_time = time.time()
             
             protect_response = self.config.galileo_platform.run_protect(
-                context, result, observe_workflow
+                context, result, galileo_logger
             )
 
+            print(protect_response)
+
             if protect_enabled and protect_response["text"] != result:
                 pii_flag["phone_number"] = "phone_number" in protect_response["metric_results"]["pii"]["value"]
                 pii_flag["email"] = "email" in protect_response["metric_results"]["pii"]["value"]
                 pii_flag["name"] = "name" in protect_response["metric_results"]["pii"]["value"]
-                # pii_flag["company"] = protect_response["metric_results"]["deutsche_bank_company_pii_0"]["value"]>0.1
+                pii_flag["company"] = protect_response["metric_results"]["deutsche_bank_company_pii_0"]["value"]>0.1
                 redacted_result = self.get_redacted_result(result, pii_flag)
+                # redacted_result = re.sub(r'<pii>(.*?)</pii>', r'<pii>REDACTED</pii>', redacted_result)
                 result = redacted_result.replace("<pii>", "<tag>").replace("</pii>", "</tag>")
-                redacted_result = re.sub(r'<pii>(.*?)</pii>', r'<pii>REDACTED</pii>', redacted_result)
 
             if hallucination_detection:
                 context_adherence_score = protect_response["metric_results"]["context_adherence_luna"]["value"]
-                # print(context_adherence_score)
-
-            # Conclude the workflow with the final result and set output
-            observe_workflow.conclude(output=result)
-            evaluate_workflow.output = result
-            self.config.galileo_platform.observe_logger.upload_workflows()
 
-            # Start evaluation in separate thread
-            self.config.galileo_platform.evaluate_run.finish(wait=True, silent=True)
-            # print(self.config.galileo_platform.evaluate_run)
+            galileo_logger.conclude(output=result)
+            galileo_logger.flush()
 
             return result, redacted_result, original_result, context_adherence_score, pii_flag
 
         except Exception as e:
-            # Log errors to Galileo Observe
-            observe_workflow.conclude(output={"error": str(e)})
-            self.config.galileo_platform.observe_logger.upload_workflows()
+            galileo_logger.conclude(output={"error": str(e)})
+            galileo_logger.flush()
             raise e
         
     def get_redacted_result(self, result, pii_flag):

backend/classes/vector_database/milvus_vector_database.py

@@ -3,7 +3,7 @@ import shutil
 from typing import List
 
 import pandas as pd
-from pymilvus import MilvusClient, connections, FieldSchema, CollectionSchema, DataType, Collection
+from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType
 import logging
 
 from backend.classes.vector_database.base_vector_database import VectorDatabaseConfig, VectorDatabase