simplify developer getting started experience (#102)

* Fixed build. Now, we have a bare bones version of the docker-compose file with only two services, archgw and archgw-model-server. Tested using CLI

* some pre-commit fixes

* fixed cargo formatting issues

* fixed model server conflict changes

---------

Co-authored-by: Salman Paracha <salmanparacha@MacBook-Pro-261.local>

model_server/app/main.py

@@ -2,7 +2,6 @@ import os
 from fastapi import FastAPI, Response, HTTPException
 from pydantic import BaseModel
 from app.load_models import (
-    load_ner_models,
     load_transformers,
     load_guard_model,
     load_zero_shot_models,
@@ -22,46 +21,18 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 transformers = load_transformers()
-ner_models = load_ner_models()
 zero_shot_models = load_zero_shot_models()
 
-config = {}
-
-if os.path.exists("/root/arch_config.yaml"):
-  with open("/root/arch_config.yaml", "r") as file:
-      config = yaml.safe_load(file)
 with open("guard_model_config.yaml") as f:
     guard_model_config = yaml.safe_load(f)
 
-if "prompt_guards" in config.keys():
-    if len(config["prompt_guards"]["input_guards"]) == 2:
-        task = "both"
-        jailbreak_hardware = "gpu" if torch.cuda.is_available() else "cpu"
-        toxic_hardware = "gpu" if torch.cuda.is_available() else "cpu"
-        toxic_model = load_guard_model(
-            guard_model_config["toxic"][jailbreak_hardware], toxic_hardware
-        )
-        jailbreak_model = load_guard_model(
-            guard_model_config["jailbreak"][toxic_hardware], jailbreak_hardware
-        )
+task = "both"
+hardware = "gpu" if torch.cuda.is_available() else "cpu"
+jailbreak_model = load_guard_model(
+    guard_model_config["jailbreak"][hardware], hardware
+)
 
-    else:
-        task = list(config["prompt_guards"]["input_guards"].keys())[0]
-
-        hardware = "gpu" if torch.cuda.is_available() else "cpu"
-        if task == "toxic":
-            toxic_model = load_guard_model(
-                guard_model_config["toxic"][hardware], hardware
-            )
-            jailbreak_model = None
-        elif task == "jailbreak":
-            jailbreak_model = load_guard_model(
-                guard_model_config["jailbreak"][hardware], hardware
-            )
-            toxic_model = None
-
-
-    guard_handler = GuardHandler(toxic_model, jailbreak_model)
+guard_handler = GuardHandler(toxic_model=None, jailbreak_model=jailbreak_model)
 
 app = FastAPI()
 
@@ -108,27 +79,6 @@ async def embedding(req: EmbeddingRequest, res: Response):
     return {"data": data, "model": req.model, "object": "list", "usage": usage}
 
 
-class NERRequest(BaseModel):
-    input: str
-    labels: list[str]
-    model: str
-
-
-@app.post("/ner")
-async def ner(req: NERRequest, res: Response):
-    if req.model not in ner_models:
-        raise HTTPException(status_code=400, detail="unknown model: " + req.model)
-
-    model = ner_models[req.model]
-    entities = model.predict_entities(req.input, req.labels)
-
-    return {
-        "data": entities,
-        "model": req.model,
-        "object": "list",
-    }
-
-
 class GuardRequest(BaseModel):
     input: str
     task: str
@@ -236,270 +186,7 @@ async def zeroshot(req: ZeroShotRequest, res: Response):
         "model": req.model,
     }
 
-
 @app.post("/v1/chat/completions")
 async def chat_completion(req: ChatMessage, res: Response):
     result = await arch_fc_chat_completion(req, res)
     return result
-
-
-'''
-*****
-Adding new functions to test the usecases - Sampreeth
-*****
-"""
-
-conn = load_sql()
-name_col = "name"
-
-
-class TopEmployees(BaseModel):
-    grouping: str
-    ranking_criteria: str
-    top_n: int
-
-
-@app.post("/top_employees")
-async def top_employees(req: TopEmployees, res: Response):
-    name_col = "name"
-    # Check if `req.ranking_criteria` is a Text object and extract its value accordingly
-    logger.info(
-        f"{'* ' * 50}\n\nCaptured Ranking Criteria: {req.ranking_criteria}\n\n{'* ' * 50}"
-    )
-
-    if req.ranking_criteria == "yoe":
-        req.ranking_criteria = "years_of_experience"
-    elif req.ranking_criteria == "rating":
-        req.ranking_criteria = "performance_score"
-
-    logger.info(
-        f"{'* ' * 50}\n\nFinal Ranking Criteria: {req.ranking_criteria}\n\n{'* ' * 50}"
-    )
-
-    query = f"""
-    SELECT {req.grouping}, {name_col}, {req.ranking_criteria}
-    FROM (
-        SELECT {req.grouping}, {name_col}, {req.ranking_criteria},
-               DENSE_RANK() OVER (PARTITION BY {req.grouping} ORDER BY {req.ranking_criteria} DESC) as emp_rank
-        FROM employees
-    ) ranked_employees
-    WHERE emp_rank <= {req.top_n};
-    """
-    result_df = pd.read_sql_query(query, conn)
-    result = result_df.to_dict(orient="records")
-    return result
-
-
-class AggregateStats(BaseModel):
-    grouping: str
-    aggregate_criteria: str
-    aggregate_type: str
-
-
-@app.post("/aggregate_stats")
-async def aggregate_stats(req: AggregateStats, res: Response):
-    logger.info(
-        f"{'* ' * 50}\n\nCaptured Aggregate Criteria: {req.aggregate_criteria}\n\n{'* ' * 50}"
-    )
-
-    if req.aggregate_criteria == "yoe":
-        req.aggregate_criteria = "years_of_experience"
-
-    logger.info(
-        f"{'* ' * 50}\n\nFinal Aggregate Criteria: {req.aggregate_criteria}\n\n{'* ' * 50}"
-    )
-
-    logger.info(
-        f"{'* ' * 50}\n\nCaptured Aggregate Type: {req.aggregate_type}\n\n{'* ' * 50}"
-    )
-    if req.aggregate_type.lower() not in ["sum", "avg", "min", "max"]:
-        if req.aggregate_type.lower() == "count":
-            req.aggregate_type = "COUNT"
-        elif req.aggregate_type.lower() == "total":
-            req.aggregate_type = "SUM"
-        elif req.aggregate_type.lower() == "average":
-            req.aggregate_type = "AVG"
-        elif req.aggregate_type.lower() == "minimum":
-            req.aggregate_type = "MIN"
-        elif req.aggregate_type.lower() == "maximum":
-            req.aggregate_type = "MAX"
-        else:
-            raise HTTPException(status_code=400, detail="Invalid aggregate type")
-
-    logger.info(
-        f"{'* ' * 50}\n\nFinal Aggregate Type: {req.aggregate_type}\n\n{'* ' * 50}"
-    )
-
-    query = f"""
-    SELECT {req.grouping}, {req.aggregate_type}({req.aggregate_criteria}) as {req.aggregate_type}_{req.aggregate_criteria}
-    FROM employees
-    GROUP BY {req.grouping};
-    """
-    result_df = pd.read_sql_query(query, conn)
-    result = result_df.to_dict(orient="records")
-    return result
-
-
-class PacketDropCorrelationRequest(BaseModel):
-    from_time: str = None  # Optional natural language timeframe
-    ifname: str = None  # Optional interface name filter
-    region: str = None  # Optional region filter
-    min_in_errors: int = None
-    max_in_errors: int = None
-    min_out_errors: int = None
-    max_out_errors: int = None
-    min_in_discards: int = None
-    max_in_discards: int = None
-    min_out_discards: int = None
-    max_out_discards: int = None
-
-
-@app.post("/interface_down_pkt_drop")
-async def interface_down_packet_drop(req: PacketDropCorrelationRequest, res: Response):
-    params, filters = load_params(req)
-
-    # Join the filters using AND
-    where_clause = " AND ".join(filters)
-    if where_clause:
-        where_clause = "AND " + where_clause
-
-    # Step 3: Query packet errors and flows from interfacestats and ts_flow
-    query = f"""
-    SELECT
-      d.switchip AS device_ip_address,
-      i.in_errors,
-      i.in_discards,
-      i.out_errors,
-      i.out_discards,
-      i.ifname,
-      t.src_addr,
-      t.dst_addr,
-      t.time AS flow_time,
-      i.time AS interface_time
-    FROM
-      device d
-    INNER JOIN
-      interfacestats i
-      ON d.device_mac_address = i.device_mac_address
-    INNER JOIN
-      ts_flow t
-      ON d.switchip = t.sampler_address
-    WHERE
-      i.time >= :from_time  -- Using the converted timestamp
-      {where_clause}
-    ORDER BY
-      i.time;
-    """
-
-    correlated_data = pd.read_sql_query(query, conn, params=params)
-
-    if correlated_data.empty:
-        default_response = {
-            "device_ip_address": "0.0.0.0",  # Placeholder IP
-            "in_errors": 0,
-            "in_discards": 0,
-            "out_errors": 0,
-            "out_discards": 0,
-            "ifname": req.ifname
-            or "unknown",  # Placeholder or interface provided in the request
-            "src_addr": "0.0.0.0",  # Placeholder source IP
-            "dst_addr": "0.0.0.0",  # Placeholder destination IP
-            "flow_time": str(
-                datetime.now(timezone.utc)
-            ),  # Current timestamp or placeholder
-            "interface_time": str(
-                datetime.now(timezone.utc)
-            ),  # Current timestamp or placeholder
-        }
-        return [default_response]
-
-    logger.info(f"Correlated Packet Drop Data: {correlated_data}")
-
-    return correlated_data.to_dict(orient='records')
-
-
-class FlowPacketErrorCorrelationRequest(BaseModel):
-    from_time: str = None  # Optional natural language timeframe
-    ifname: str = None  # Optional interface name filter
-    region: str = None  # Optional region filter
-    min_in_errors: int = None
-    max_in_errors: int = None
-    min_out_errors: int = None
-    max_out_errors: int = None
-    min_in_discards: int = None
-    max_in_discards: int = None
-    min_out_discards: int = None
-    max_out_discards: int = None
-
-
-@app.post("/packet_errors_impact_flow")
-async def packet_errors_impact_flow(
-    req: FlowPacketErrorCorrelationRequest, res: Response
-):
-    params, filters = load_params(req)
-
-    # Join the filters using AND
-    where_clause = " AND ".join(filters)
-    if where_clause:
-        where_clause = "AND " + where_clause
-
-    # Step 3: Query the packet errors and flows, correlating by timestamps
-    query = f"""
-    SELECT
-      d.switchip AS device_ip_address,
-      i.in_errors,
-      i.in_discards,
-      i.out_errors,
-      i.out_discards,
-      i.ifname,
-      t.src_addr,
-      t.dst_addr,
-      t.src_port,
-      t.dst_port,
-      t.packets,
-      t.time AS flow_time,
-      i.time AS error_time
-    FROM
-      device d
-    INNER JOIN
-      interfacestats i
-      ON d.device_mac_address = i.device_mac_address
-    INNER JOIN
-      ts_flow t
-      ON d.switchip = t.sampler_address
-    WHERE
-      i.time >= :from_time
-      AND ABS(strftime('%s', t.time) - strftime('%s', i.time)) <= 300  -- Correlate within 5 minutes
-      {where_clause}
-    ORDER BY
-      i.time;
-    """
-
-    correlated_data = pd.read_sql_query(query, conn, params=params)
-
-    if correlated_data.empty:
-        default_response = {
-            "device_ip_address": "0.0.0.0",  # Placeholder IP
-            "in_errors": 0,
-            "in_discards": 0,
-            "out_errors": 0,
-            "out_discards": 0,
-            "ifname": req.ifname
-            or "unknown",  # Placeholder or interface provided in the request
-            "src_addr": "0.0.0.0",  # Placeholder source IP
-            "dst_addr": "0.0.0.0",  # Placeholder destination IP
-            "src_port": 0,
-            "dst_port": 0,
-            "packets": 0,
-            "flow_time": str(
-                datetime.now(timezone.utc)
-            ),  # Current timestamp or placeholder
-            "error_time": str(
-                datetime.now(timezone.utc)
-            ),  # Current timestamp or placeholder
-        }
-        return [default_response]
-
-    # Return the correlated data if found
-    return correlated_data.to_dict(orient="records")
-'''