feat: add openai for client

langchain_mcp_client.py

@@ -39,19 +39,19 @@ async def lc_mcp_exec(request: str, history=None) -> Tuple[str, list]:
         server_params = get_server_params()
 
         # Initialize the LLM for OpenAI
-        # llm = init_chat_model(
-        #     model_provider=os.environ["OPENAI_MODEL_PROVIDER"],
-        #     model=os.environ["OPENAI_MODEL"],
-        #     api_key=os.environ["OPENAI_API_KEY"]
-        # )
-
-        # Initialize the LLM for Gemini
         llm = init_chat_model(
-            model_provider=os.environ["GEMINI_MODEL_PROVIDER"],
-            model=os.environ["GEMINI_MODEL"],
-            api_key=os.environ["GEMINI_API_KEY"]
+            model_provider=os.environ["OPENAI_MODEL_PROVIDER"],
+            model=os.environ["OPENAI_MODEL"],
+            api_key=os.environ["OPENAI_API_KEY"]
         )
 
+        # Initialize the LLM for Gemini
+        # llm = init_chat_model(
+        #     model_provider=os.environ["GEMINI_MODEL_PROVIDER"],
+        #     model=os.environ["GEMINI_MODEL"],
+        #     api_key=os.environ["GEMINI_API_KEY"]
+        # )
+
         # Initialize the MCP client
         async with stdio_client(server_params) as (read, write):
             async with ClientSession(read, write) as session:

postgre_mcp_server.py

@@ -2,14 +2,14 @@ import os
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from typing import Optional, AsyncIterator
-import asyncpg
-from dotenv import load_dotenv
+import asyncpg 
 from mcp.server.fastmcp import FastMCP, Context
 from pydantic import Field
 import pandasai as pai
 import matplotlib as plt
 import pandas as pd
 import logging
+from pandasai_openai import OpenAI
 
 
 # Constants
@@ -21,8 +21,8 @@ DEFAULT_QUERY_LIMIT = 100
 # get logger
 logger = logging.getLogger(__name__)
 
-# Load environment variables
-load_dotenv()
+pllm = OpenAI(api_token=os.environ["OPENAI_API_KEY"])
+pai.config.set({"llm": pllm})
 
 # Define our own PromptMessage class if the MCP one isn't available
 @dataclass
@@ -37,16 +37,11 @@ class DbContext:
     pool: asyncpg.Pool
     schema: str
 
-# Load environment variables
-load_dotenv()
-
 
 # Database connection lifecycle manager
 @asynccontextmanager
 async def db_lifespan(server: FastMCP) -> AsyncIterator[DbContext]:
     """Manage database connection lifecycle"""
-    # Initialize DB connection from environment variables
-    load_dotenv()
     dsn = os.environ["DB_URL"]
     schema = os.environ["DB_SCHEMA"]
 
@@ -453,8 +448,7 @@ async def generate_analytical_query(table_name: str) -> list[PromptMessage]:
     """ Generate analytical queries for a table
     Args:
         table_name: The name of the table to generate analytical queries for
-    """
-    load_dotenv()
+    """ 
     schema = os.environ["DB_SCHEMA"]
     try:
         async with db_lifespan(mcp) as db_ctx:
@@ -627,11 +621,7 @@ async def visualize_results(json_data: dict, vis_prompt: str) -> str:
     """
     try:
         # Debug prints to see what's being received
-        from pandasai_openai import OpenAI
 
-        pllm = OpenAI(api_token=os.environ["OPENAI_API_KEY"])
-        pai.config.set({"llm": pllm})
-        
         # Convert JSON to DataFrame
         df = pd.DataFrame(json_data["data"], columns=json_data["columns"])
         
@@ -642,8 +632,6 @@ async def visualize_results(json_data: dict, vis_prompt: str) -> str:
         # Initialize PandasAI
         df_ai = pai.DataFrame(df)
 
-        # Load api key
-        load_dotenv()
         api_key = os.environ["PANDAS_KEY"]
         pai.api_key.set(api_key)