Deva22467/govt-scheme-openenv

1import os2import json3import argparse4from typing import Dict, Any5from env import SchemeEnv6from models import Observation, Action7 8class RuleBasedBaselineAgent:9    """Robust rule-based logic natively navigating across explicit trajectory limits"""10    def act(self, state: Observation) -> Action:11        occupation = getattr(state, "occupation", "").lower()12        age = getattr(state, "age", 0)13        category = getattr(state, "category", "")14        stage = getattr(state, "stage", 0)15        16        schemes = []17        best = ""18        reason = ""19        20        if "farmer" in occupation:21            schemes = ["PM Kisan", "Crop Insurance"]22            best = "PM Kisan"23            reason = "Farmer from low income group needing direct income support."24        elif "student" in occupation:25            schemes = ["Post Matric Scholarship OBC", "Sukanya Samriddhi Yojana"]26            best = "Post Matric Scholarship OBC" if category == "OBC" else "Sukanya Samriddhi Yojana"27            reason = "Student pursuing higher education."28        elif "startup" in occupation:29            schemes = ["Startup India Seed Fund"]30            best = "Startup India Seed Fund"31            reason = "Women entrepreneur targeting seed funding."32        elif age >= 60:33            schemes = ["IGNOAPS", "MGNREGA"]34            best = "IGNOAPS"35            reason = "Elderly citizen requires passive pension support."36        elif "unemployed" in occupation:37            schemes = ["PM Kaushal Vikas Yojana", "NDIS"]38            best = "PM Kaushal Vikas Yojana"39            reason = "Unemployed youth needs skills."40        else:41            schemes = ["MGNREGA", "PMAY-G", "PM Van Dhan Yojana"]42            best = "PMAY-G"43            reason = "Rural laborer primarily benefits from housing scheme."44            45        if stage == 0:46            return Action(schemes=schemes)47        elif stage == 1:48            return Action(best=best)49        elif stage == 2:50            return Action(reasoning=reason)51        else:52            return Action()53 54class OpenAIBaselineAgent:55    """True OpenEnv Model Inference Execution referencing default OPENAI_API_KEY variables"""56    def __init__(self):57        try:58            from openai import OpenAI59            self.client = OpenAI() # Inherits OPENAI_API_KEY directly from environment properties60        except ImportError:61            raise ImportError("Please install openai: pip install openai")62 63    def act(self, state: Observation) -> Action:64        stage = getattr(state, "stage", 0)65        if stage == 0:66            format_str = '{"schemes": ["Scheme A", "Scheme B"]}'67        elif stage == 1:68            format_str = '{"best": "Best scheme name"}'69        else:70            format_str = '{"reasoning": "Rationale description"}'71 72        prompt = f"""73        Citizen profile:74        {state.model_dump_json()}75        76        Stage {stage} trajectory required.77        Return explicitly valid JSON matching:78        {format_str}79        """80        81        try:82            response = self.client.chat.completions.create(83                model="gpt-4o-mini",84                response_format={"type": "json_object"},85                messages=[{"role":"user","content":prompt}]86            )87            raw_json = response.choices[0].message.content88            action_data = json.loads(raw_json)89            return Action(**action_data)90        except Exception as e:91            print(f"OpenAI Integration Error during Stage {stage}: {e}")92            return Action()93 94def demo_agent(agent_type="rule", num_episodes=1):95    print(f"Demostrating {agent_type} agent over {num_episodes} episode...")96    env = SchemeEnv()97    98    if agent_type == "openai" and os.environ.get("OPENAI_API_KEY"):99        agent = OpenAIBaselineAgent()100    else:101        if agent_type == "openai":102            print("Warning: Missing OPENAI_API_KEY. Falling back to baseline rule-based routing.")103        agent = RuleBasedBaselineAgent()104        105    for i in range(num_episodes):106        obs = env.reset()107        done = False108        print(f"\n--- Episode {i+1} ---")109        while not done:110            action = agent.act(obs)111            obs, reward, done, info = env.step(action)112            print(f"Stage: {info['stage']}, Action: {action.model_dump_json()}, Yield: {reward:.2f}")113 114if __name__ == "__main__":115    parser = argparse.ArgumentParser()116    parser.add_argument("--agent", type=str, default="rule", choices=["rule", "openai"])117    args = parser.parse_args()118    demo_agent(args.agent, num_episodes=1)119