integrate gemini and claude as ai providers to mcp endpoint

Author: smutyala1at

apps/meshjs-rag/app/api/v1/ask_mesh_ai.py

@@ -48,7 +48,11 @@ async def ask_mesh_ai(body: ChatCompletionRequest, credentials: HTTPAuthorizatio
   if openai_api_key is None:
     raise ValueError("OpenAI api key is missing")
 
-  openai_service = OpenAIService(openai_api_key)
+  openai_service = OpenAIService(
+    embedding_api_key=openai_api_key,
+    completion_api_key=openai_api_key,
+    completion_model="gpt-4o-mini"
+  )
 
   try:
     question = body.messages[-1].content
@@ -71,35 +75,47 @@ async def ask_mesh_ai(body: ChatCompletionRequest, credentials: HTTPAuthorizatio
 
 ###########################################################################################################
 @router.post("/mcp")
-async def ask_mesh_ai(body: MCPRequestBody, authorization: str = Header(None), supabase: AsyncClient = Depends(get_db_client)):
+async def mesh_mcp(body: MCPRequestBody, authorization: str = Header(None), supabase: AsyncClient = Depends(get_db_client)):
 
   if not authorization or not authorization.startswith("Bearer"):
     print("error")
     raise HTTPException(
       status_code=status.HTTP_401_UNAUTHORIZED,
       detail="You are not authorized"
     )
+  
+  embedding_api_key = os.getenv("OPENAI_KEY") or None
+  if embedding_api_key is None:
+    raise ValueError("Embedding api key is missing")
 
   try:
-    OPENAI_KEY = authorization.split(" ")[-1]
-    openai_service = OpenAIService(OPENAI_KEY)
-
+    completion_api_key = authorization.split(" ")[-1]
     question = body.query
     model = body.model
 
+    if model.startswith("gemini"):
+      base_url = "https://generativelanguage.googleapis.com/v1beta/openai/"
+    elif model.startswith("claude"):
+      base_url = "https://api.anthropic.com/v1/"
+
+    openai_service = OpenAIService(
+      embedding_api_key=embedding_api_key,
+      completion_api_key=completion_api_key,
+      completion_model=model,
+      base_url=base_url
+    )
+
     embedded_query = await openai_service.embed_query(question)
     context = await get_context(embedded_query, supabase)
-    response = await openai_service.get_mcp_answer(question=question, context=context, model=model)
+    response = await openai_service.get_mcp_answer(question=question, context=context)
     return response
 
   except (openai.APIError, openai.AuthenticationError, openai.RateLimitError) as e:
-    print(e)
     raise HTTPException(
       status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
       detail=f"An OpenAI API error occurred: {e}"
     )
   except Exception as e:
-    print(e)
     raise HTTPException(
       status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
       detail=f"An unexpected error occurred: {e}"

apps/meshjs-rag/app/services/openai.py

@@ -36,13 +36,18 @@
 """
 
 class OpenAIService:
-  def __init__(self, openai_api_key):
-    self.client = AsyncOpenAI(api_key=openai_api_key)
+  def __init__(self, embedding_api_key: str, completion_api_key: str, completion_model: str, base_url: str = None):
+    self.embedding_client = AsyncOpenAI(api_key=embedding_api_key)
+    self.completion_client = AsyncOpenAI(
+      api_key=completion_api_key,
+      base_url=base_url
+    )
+    self.model = completion_model
 
   @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
-  async def _chat(self, messages, model="gpt-4o-mini", temperature=0.0, max_tokens=None, prompt_cache_key=None, stream: bool = False):
+  async def _chat(self, messages, temperature=0.0, max_tokens=None, prompt_cache_key=None, stream: bool = False):
     kwargs = {
-      "model": model,
+      "model": self.model,
       "messages": messages,
       "temperature": temperature,
       "stream": stream
@@ -53,7 +58,7 @@ async def _chat(self, messages, model="gpt-4o-mini", temperature=0.0, max_tokens
     if max_tokens:
       kwargs["max_tokens"] = max_tokens
 
-    return await self.client.chat.completions.create(**kwargs)
+    return await self.completion_client.chat.completions.create(**kwargs)
 
   async def situate_context(self, doc: str, chunk: str, cache_key: str) -> str:
     messages = [
@@ -72,7 +77,7 @@ async def situate_context(self, doc: str, chunk: str, cache_key: str) -> str:
 
   @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6), reraise=True)
   async def get_batch_embeddings(self, texts: List[str]) -> List[List[float]]:
-    response = await self.client.embeddings.create(
+    response = await self.embedding_client.embeddings.create(
       model="text-embedding-3-small",
       input=texts,
       encoding_format="float"
@@ -82,7 +87,7 @@ async def get_batch_embeddings(self, texts: List[str]) -> List[List[float]]:
 
   @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6), reraise=True)
   async def embed_query(self, text: str) -> List[float]:
-    response = await self.client.embeddings.create(
+    response = await self.embedding_client.embeddings.create(
       model="text-embedding-3-small",
       input=text,
       encoding_format="float"
@@ -91,7 +96,7 @@ async def embed_query(self, text: str) -> List[float]:
     return response.data[0].embedding
 
   @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6), reraise=True)
-  async def get_answer(self, question: str, context: str, model="gpt-4o-mini"):
+  async def get_answer(self, question: str, context: str):
     messages = [
       {
         "role": "system",
@@ -103,15 +108,15 @@ async def get_answer(self, question: str, context: str, model="gpt-4o-mini"):
       }
     ]
 
-    stream = await self._chat(messages=messages, stream=True, model=model)
+    stream = await self._chat(messages=messages, stream=True)
 
     async for chunk in stream:
       yield f"data: {json.dumps(chunk.model_dump())}\n\n"
 
     yield "data: [DONE]\n\n"
 
   @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6), reraise=True)
-  async def get_mcp_answer(self, question: str, context: str, model="gpt-4o-mini"):
+  async def get_mcp_answer(self, question: str, context: str):
     messages = [
       {
         "role": "system",
@@ -123,5 +128,5 @@ async def get_mcp_answer(self, question: str, context: str, model="gpt-4o-mini")
       }
     ]
 
-    response = await self._chat(messages=messages, model=model)
+    response = await self._chat(messages=messages)
     return response.choices[0].message.content

apps/meshjs-rag/app/utils/process_chunks.py

@@ -13,7 +13,11 @@
 if openai_api_key is None:
   raise ValueError("OpenAI api key is missing")
 
-openai_service = OpenAIService(openai_api_key=openai_api_key)
+openai_service = OpenAIService(
+  embedding_api_key=openai_api_key,
+  completion_api_key=openai_api_key,
+  completion_model="gpt-4o-mini"
+)
 
 async def process_chunks_and_update_db(
     chunks: List[str],