Merge pull request #104 from oramasearch/feat/auth-grpc

feat: adds auth to gRPC

src/ai_server/src/actions/main.py

@@ -12,12 +12,17 @@ class Actions:
     def __init__(self, config: OramaAIConfig):
         self.config = config
 
-    def call_oramacore_search(self, collection_id: str, query: Dict) -> any:
+    def call_oramacore_search(self, collection_id: str, query: Dict, api_key: str) -> any:
         body = json.dumps(query)
-        url = f"http://{self.config.rust_server_host}:{self.config.rust_server_port}/v1/{collection_id}/actions/execute"
-
+        url = f"http://{self.config.rust_server_host}:{self.config.rust_server_port}/v1/{collection_id}/actions/execute?api-key={api_key}"
         headers = {"Content-Type": "application/json; charset=utf-8"}
-        resp = requests.post(url=url, json={"context": body, "name": "search"}, headers=headers)
-        as_json = json.loads(resp.text)
 
-        return as_json["hits"]
+        try:
+            resp = requests.post(url=url, json={"context": body, "name": "search"}, headers=headers)
+            as_json = json.loads(resp.text)
+            return as_json["hits"]
+
+        except Exception as e:
+            print(e)
+            logger.error(f"Error calling oramacore search: {e}")
+            return None

src/ai_server/src/actions/party_planner.py

@@ -56,34 +56,32 @@ def _create_step(self, action: Dict[str, str]) -> Step:
             should_stream=step_config["stream"],
         )
 
-    def _execute_orama_search(self, collection_id: str, input: str):
-
-        print(json.dumps(self.executed_steps, indent=2))
+    def _execute_orama_search(self, collection_id: str, input: str, api_key: str) -> List[Dict[str, Any]]:
+        queries = []
 
         for step in self.executed_steps:
+            if step.action == "GENERATE_QUERIES":
+                queries = step.result
+            elif step.action == "OPTIMIZE_QUERY":
+                queries = [step.result]
+            else:
+                queries = [input]
 
-            if step.action == "OPTIMIZE_QUERY":
-                return self.act.call_oramacore_search(
-                    collection_id, {"term": step.result, "mode": "hybrid", "limit": 5}
-                )
-
-            elif step.action == "GENERATE_QUERIES":
-                results = []
-
-                for query in step.result:
-                    res = self.act.call_oramacore_search(collection_id, {"term": query, "mode": "hybrid", "limit": 3})
-                    results.append(res)
+        results = []
+        limit = 3 if len(queries) > 1 else 5
 
-                return results
+        for query in queries:
+            full_query = {"term": query, "mode": "hybrid", "limit": limit}
+            res = self.act.call_oramacore_search(collection_id=collection_id, query=full_query, api_key=api_key)
+            results.append(res)
 
-            else:
-                return self.act.call_oramacore_search(collection_id, {"term": input, "mode": "hybrid", "limit": 5})
+        return results
 
-    def _handle_orama_step(self, step: Step, collection_id: str, input: str) -> str:
+    def _handle_orama_step(self, step: Step, collection_id: str, input: str, api_key: str) -> str:
         """Handle Orama-specific steps."""
         if step.name == "PERFORM_ORAMA_SEARCH":
             try:
-                result = self._execute_orama_search(collection_id, input)
+                result = self._execute_orama_search(collection_id=collection_id, input=input, api_key=api_key)
 
                 return json.dumps(result) if isinstance(result, dict) else str(result)
             except Exception as e:
@@ -107,7 +105,7 @@ def _handle_streaming_step(self, step: Step, input: str, history: List[Any]) ->
         history.append({"role": "assistant", "content": accumulated_result})
         yield accumulated_result, history
 
-    def run(self, collection_id: str, input: str, history: List[Any]) -> Iterator[str]:
+    def run(self, collection_id: str, input: str, history: List[Any], api_key: str) -> Iterator[str]:
         action_plan = self._get_action_plan(history, input)
         message = Message("ACTION_PLAN", action_plan)
         self.executed_steps.append(message)
@@ -117,20 +115,24 @@ def run(self, collection_id: str, input: str, history: List[Any]) -> Iterator[st
         for action in action_plan:
             step = self._create_step(action)
 
+            # Handle Orama-specific steps first
             if step.is_orama_step:
-                result = self._handle_orama_step(step, collection_id, input)
+                result = self._handle_orama_step(step=step, collection_id=collection_id, input=input, api_key=api_key)
                 message = Message(step.name, result)
                 self.executed_steps.append(message)
                 yield message.to_json()
                 continue
 
+            # Handle non-streaming and streaming steps
             if not step.should_stream:
                 result, history = self._handle_non_streaming_step(step, input, [])
 
                 message = Message(step.name, result)
                 self.executed_steps.append(message)
 
                 yield message.to_json()
+
+            # Handle streaming steps
             else:
                 step_result_acc = Message(step.name, "")
                 for result, updated_history in self._handle_streaming_step(step, input, []):

src/ai_server/src/grpc/server.py

@@ -93,6 +93,9 @@ def ChatStream(self, request, context):
             context.set_details(f"Error in chat stream: {str(e)}")
 
     def PlannedAnswer(self, request, context):
+        metadata = dict(context.invocation_metadata())
+        api_key = metadata.get("x-api-key")
+
         try:
             history = (
                 [
@@ -103,7 +106,9 @@ def PlannedAnswer(self, request, context):
                 else []
             )
 
-            for message in self.party_planner.run(request.collection_id, request.input, history):
+            for message in self.party_planner.run(
+                collection_id=request.collection_id, input=request.input, history=history, api_key=api_key
+            ):
                 yield PlannedAnswerResponse(data=message, finished=False)
 
             yield PlannedAnswerResponse(data="", finished=True)
@@ -115,10 +120,29 @@ def PlannedAnswer(self, request, context):
             return PlannedAnswerResponse()
 
 
+class AuthInterceptor(grpc.ServerInterceptor):
+    def intercept_service(self, continuation, handler_call_details):
+
+        # Health check and embeddings won't require authentication.
+        # This server should never be exposed to the public and it's meant for internal use only.
+        allowed_methods = ["CheckHealth", "GetEmbedding", "ServerReflection"]
+        if any(x in handler_call_details.method for x in allowed_methods):
+            return continuation(handler_call_details)
+
+        # The current gRPC server is a proxy for the Rust server, which requires an API key.
+        # There's no API key validation in the Python server, so we just check if the API key is present.
+        metadata = dict(handler_call_details.invocation_metadata)
+        if "x-api-key" not in metadata:
+            return grpc.unary_unary_rpc_method_handler(
+                lambda req, ctx: ctx.abort(grpc.StatusCode.UNAUTHENTICATED, "Missing API key")
+            )
+        return continuation(handler_call_details)
+
+
 def serve(config, embeddings_service, models_manager):
     logger = logging.getLogger(__name__)
     logger.info(f"Starting gRPC server on port {config.port}")
-    server = grpc.server(ThreadPoolExecutor(max_workers=10))
+    server = grpc.server(ThreadPoolExecutor(max_workers=10), interceptors=[AuthInterceptor()])
     logger.info("gRPC server created")
 
     llm_service = LLMService(embeddings_service, models_manager, config)

src/ai_server/test.sh

@@ -8,16 +8,16 @@ grpcurl -plaintext localhost:50051 orama_ai_service.LLMService/CheckHealth
 grpcurl -plaintext -d '{ "model": "MultilingualE5Small", "input": ["The quick brown fox jumps over the lazy dog"], "intent": "passage" }' localhost:50051 orama_ai_service.LLMService/GetEmbedding
 
 # Chat (non-streaming)
-grpcurl -plaintext -d '{ "model": "google_query_translator", "prompt": "I am installing my Ryzen 9 9900X and I fear I bent some pins. What should I do?" }' localhost:50051 orama_ai_service.LLMService/Chat
+grpcurl -plaintext -H 'x-api-key: read_api_key' -d '{ "model": "google_query_translator", "prompt": "I am installing my Ryzen 9 9900X and I fear I bent some pins. What should I do?" }' localhost:50051 orama_ai_service.LLMService/Chat
 
 # Chat (streaming)
-grpcurl -plaintext -d '{ "model": "google_query_translator", "prompt": "I am installing my Ryzen 9 9900X and I fear I bent some pins. What should I do?" }' localhost:50051 orama_ai_service.LLMService/ChatStream
+grpcurl -plaintext -H 'x-api-key: read_api_key' -d '{ "model": "google_query_translator", "prompt": "I am installing my Ryzen 9 9900X and I fear I bent some pins. What should I do?" }' localhost:50051 orama_ai_service.LLMService/ChatStream
 
 # Answer with context (streaming)
-grpcurl -plaintext -d '{ "model": "answer", "prompt": "What do you know about beavers?", "conversation": { "messages": [ { "role": "USER", "content": "Beavers (genus Castor) are large, semiaquatic rodents of the Northern Hemisphere. There are two existing species: the North American beaver (Castor canadensis) and the Eurasian beaver (C. fiber). Beavers are the second-largest living rodents, after capybaras, weighing up to 50 kg (110 lb). They have stout bodies with large heads, long chisel-like incisors, brown or gray fur, hand-like front feet, webbed back feet, and tails that are flat and scaly. The two species differ in skull and tail shape and fur color. Beavers can be found in a number of freshwater habitats, such as rivers, streams, lakes and ponds. They are herbivorous, consuming tree bark, aquatic plants, grasses and sedges." } ] } }' localhost:50051 orama_ai_service.LLMService/ChatStream
+grpcurl -plaintext -H 'x-api-key: read_api_key' -d '{ "model": "answer", "prompt": "What do you know about beavers?", "conversation": { "messages": [ { "role": "USER", "content": "Beavers (genus Castor) are large, semiaquatic rodents of the Northern Hemisphere. There are two existing species: the North American beaver (Castor canadensis) and the Eurasian beaver (C. fiber). Beavers are the second-largest living rodents, after capybaras, weighing up to 50 kg (110 lb). They have stout bodies with large heads, long chisel-like incisors, brown or gray fur, hand-like front feet, webbed back feet, and tails that are flat and scaly. The two species differ in skull and tail shape and fur color. Beavers can be found in a number of freshwater habitats, such as rivers, streams, lakes and ponds. They are herbivorous, consuming tree bark, aquatic plants, grasses and sedges." } ] } }' localhost:50051 orama_ai_service.LLMService/ChatStream
 
 # Answer with context (non-streaming)
-grpcurl -plaintext -d '{ "model": "answer", "prompt": "What do you know about beavers?", "conversation": { "messages": [ { "role": "USER", "content": "Beavers (genus Castor) are large, semiaquatic rodents of the Northern Hemisphere. There are two existing species: the North American beaver (Castor canadensis) and the Eurasian beaver (C. fiber). Beavers are the second-largest living rodents, after capybaras, weighing up to 50 kg (110 lb). They have stout bodies with large heads, long chisel-like incisors, brown or gray fur, hand-like front feet, webbed back feet, and tails that are flat and scaly. The two species differ in skull and tail shape and fur color. Beavers can be found in a number of freshwater habitats, such as rivers, streams, lakes and ponds. They are herbivorous, consuming tree bark, aquatic plants, grasses and sedges." } ] } }' localhost:50051 orama_ai_service.LLMService/Chat
+grpcurl -plaintext -H 'x-api-key: read_api_key' -d '{ "model": "answer", "prompt": "What do you know about beavers?", "conversation": { "messages": [ { "role": "USER", "content": "Beavers (genus Castor) are large, semiaquatic rodents of the Northern Hemisphere. There are two existing species: the North American beaver (Castor canadensis) and the Eurasian beaver (C. fiber). Beavers are the second-largest living rodents, after capybaras, weighing up to 50 kg (110 lb). They have stout bodies with large heads, long chisel-like incisors, brown or gray fur, hand-like front feet, webbed back feet, and tails that are flat and scaly. The two species differ in skull and tail shape and fur color. Beavers can be found in a number of freshwater habitats, such as rivers, streams, lakes and ponds. They are herbivorous, consuming tree bark, aquatic plants, grasses and sedges." } ] } }' localhost:50051 orama_ai_service.LLMService/Chat
 
 # Party Planner (streaming)
-grpcurl -plaintext -d '{ "input": "I just started playing basketball and I want a good pair of shoes. Can you help me choose one?", "collection_id": "nike-data", "conversation": [] }' localhost:50051 orama_ai_service.LLMService/PlannedAnswer
+grpcurl -plaintext -H 'x-api-key: read_api_key' -d '{ "input": "I just started playing basketball and I want a good pair of shoes. Can you help me choose one?", "collection_id": "nike-data", "conversation": [] }' localhost:50051 orama_ai_service.LLMService/PlannedAnswer