fix(typing): strictly type openrouter, simplify imports and routing logic

Author: dineshpinto

src/flare_ai_consensus/consensus/__init__.py

@@ -0,0 +1,12 @@
+from .aggregator import async_centralized_llm_aggregator, centralized_llm_aggregator
+from .config import AggregatorConfig, ConsensusConfig, ModelConfig
+from .consensus import send_round
+
+__all__ = [
+    "AggregatorConfig",
+    "ConsensusConfig",
+    "ModelConfig",
+    "async_centralized_llm_aggregator",
+    "centralized_llm_aggregator",
+    "send_round",
+]

src/flare_ai_consensus/consensus/aggregator.py

@@ -1,5 +1,10 @@
 from flare_ai_consensus.consensus.config import AggregatorConfig
-from flare_ai_consensus.router.client import AsyncOpenRouterClient, OpenRouterClient
+from flare_ai_consensus.router import (
+    AsyncOpenRouterProvider,
+    ChatRequest,
+    Message,
+    OpenRouterProvider,
+)
 
 
 def concatenate_aggregator(responses: dict[str, str]) -> str:
@@ -13,20 +18,20 @@ def concatenate_aggregator(responses: dict[str, str]) -> str:
 
 
 def centralized_llm_aggregator(
-    client: OpenRouterClient,
+    provider: OpenRouterProvider,
     aggregator_config: AggregatorConfig,
     aggregated_responses: dict[str, str],
 ) -> str:
     """Use a centralized LLM  to combine responses.
 
-    :param client: An OpenRouterClient instance.
+    :param provider: An OpenRouterProvider instance.
     :param aggregator_config: An instance of AggregatorConfig.
     :param aggregated_responses: A string containing aggregated
         responses from individual models.
     :return: The aggregator's combined response.
     """
     # Build the message list.
-    messages = []
+    messages: list[Message] = []
     messages.extend(aggregator_config.context)
 
     # Add a system message with the aggregated responses.
@@ -38,27 +43,27 @@ def centralized_llm_aggregator(
     # Add the aggregator prompt
     messages.extend(aggregator_config.prompt)
 
-    payload = {
+    payload: ChatRequest = {
         "model": aggregator_config.model.model_id,
         "messages": messages,
         "max_tokens": aggregator_config.model.max_tokens,
         "temperature": aggregator_config.model.temperature,
     }
 
     # Get aggregated response from the centralized LLM
-    response = client.send_chat_completion(payload)
+    response = provider.send_chat_completion(payload)
     return response.get("choices", [])[0].get("message", {}).get("content", "")
 
 
 async def async_centralized_llm_aggregator(
-    client: AsyncOpenRouterClient,
+    provider: AsyncOpenRouterProvider,
     aggregator_config: AggregatorConfig,
     aggregated_responses: dict[str, str],
 ) -> str:
     """
-    Use a centralized LLM (via an async client) to combine responses.
+    Use a centralized LLM (via an async provider) to combine responses.
 
-    :param client: An asynchronous OpenRouter client.
+    :param provider: An asynchronous OpenRouterProvider.
     :param aggregator_config: An instance of AggregatorConfig.
     :param aggregated_responses: A string containing aggregated
         responses from individual models.
@@ -71,12 +76,12 @@ async def async_centralized_llm_aggregator(
     )
     messages.extend(aggregator_config.prompt)
 
-    payload = {
+    payload: ChatRequest = {
         "model": aggregator_config.model.model_id,
         "messages": messages,
         "max_tokens": aggregator_config.model.max_tokens,
         "temperature": aggregator_config.model.temperature,
     }
 
-    response = await client.send_chat_completion(payload)
+    response = await provider.send_chat_completion(payload)
     return response.get("choices", [])[0].get("message", {}).get("content", "")

src/flare_ai_consensus/consensus/config.py

@@ -1,9 +1,12 @@
 from dataclasses import dataclass
+from typing import Literal
+
+from flare_ai_consensus.router import Message
 
 
 @dataclass(frozen=True)
 class ModelConfig:
-    model_id: str | None = None
+    model_id: str
     max_tokens: int = 50
     temperature: float = 0.7
 
@@ -12,18 +15,18 @@ class ModelConfig:
 class AggregatorConfig:
     model: ModelConfig
     approach: str
-    context: list[dict]
-    prompt: list[dict]
+    context: list[Message]
+    prompt: list[Message]
 
 
 @dataclass(frozen=True)
 class ConsensusConfig:
     models: list[ModelConfig]
     aggregator_config: AggregatorConfig
-    initial_prompt: list[dict]
+    initial_prompt: list[Message]
     improvement_prompt: str
     iterations: int
-    aggregated_prompt_type: str
+    aggregated_prompt_type: Literal["user", "assistant", "system"]
 
     @staticmethod
     def load_parameters(json_data: dict) -> "ConsensusConfig":

src/flare_ai_consensus/consensus/consensus.py

@@ -3,15 +3,15 @@
 import structlog
 
 from flare_ai_consensus.consensus.config import ConsensusConfig, ModelConfig
-from flare_ai_consensus.router.client import AsyncOpenRouterClient
-from flare_ai_consensus.utils.parser import parse_chat_response
+from flare_ai_consensus.router import AsyncOpenRouterProvider, ChatRequest, Message
+from flare_ai_consensus.utils import parse_chat_response
 
 logger = structlog.get_logger(__name__)
 
 
 def build_improvement_conversation(
     consensus_config: ConsensusConfig, aggregated_response: str
-) -> list:
+) -> list[Message]:
     """Build an updated conversation using the consensus configuration.
 
     :param consensus_config: An instance of ConsensusConfig.
@@ -36,15 +36,15 @@ def build_improvement_conversation(
 
 
 async def get_response_for_model(
-    client: AsyncOpenRouterClient,
+    provider: AsyncOpenRouterProvider,
     consensus_config: ConsensusConfig,
     model: ModelConfig,
     aggregated_response: str | None,
 ) -> tuple[str | None, str]:
     """
     Asynchronously sends a chat completion request for a given model.
 
-    :param client: An instance of an asynchronous OpenRouter client.
+    :param provider: An instance of an asynchronous OpenRouter provider.
     :param consensus_config: An instance of ConsensusConfig.
     :param aggregated_response: The aggregated consensus response
         from the previous round (or None).
@@ -62,34 +62,34 @@ async def get_response_for_model(
         )
         logger.info("sending improvement prompt", model_id=model.model_id)
 
-    payload = {
+    payload: ChatRequest = {
         "model": model.model_id,
         "messages": conversation,
         "max_tokens": model.max_tokens,
         "temperature": model.temperature,
     }
-    response = await client.send_chat_completion(payload)
+    response = await provider.send_chat_completion(payload)
     text = parse_chat_response(response)
     logger.info("new response", model_id=model.model_id, response=text)
     return model.model_id, text
 
 
 async def send_round(
-    client: AsyncOpenRouterClient,
+    provider: AsyncOpenRouterProvider,
     consensus_config: ConsensusConfig,
     aggregated_response: str | None = None,
 ) -> dict:
     """
     Asynchronously sends a round of chat completion requests for all models.
 
-    :param client: An instance of an asynchronous OpenRouter client.
+    :param provider: An instance of an asynchronous OpenRouter provider.
     :param consensus_config: An instance of ConsensusConfig.
     :param aggregated_response: The aggregated consensus response from the
         previous round (or None).
     :return: A dictionary mapping model IDs to their response texts.
     """
     tasks = [
-        get_response_for_model(client, consensus_config, model, aggregated_response)
+        get_response_for_model(provider, consensus_config, model, aggregated_response)
         for model in consensus_config.models
     ]
     results = await asyncio.gather(*tasks)

src/flare_ai_consensus/main.py

@@ -3,35 +3,35 @@
 import structlog
 
 from flare_ai_consensus.config import config
-from flare_ai_consensus.consensus import aggregator, consensus
-from flare_ai_consensus.consensus.config import ConsensusConfig
-from flare_ai_consensus.router.client import AsyncOpenRouterClient
-from flare_ai_consensus.utils import (
-    loader,
-    saver,
+from flare_ai_consensus.consensus import (
+    ConsensusConfig,
+    async_centralized_llm_aggregator,
+    send_round,
 )
+from flare_ai_consensus.router import AsyncOpenRouterProvider
+from flare_ai_consensus.utils import load_json, save_json
 
 logger = structlog.get_logger(__name__)
 
 
 async def run_consensus(
-    client: AsyncOpenRouterClient,
+    provider: AsyncOpenRouterProvider,
     consensus_config: ConsensusConfig,
 ) -> None:
     """
     Asynchronously runs the consensus learning loop.
 
-    :param client: An instance of a synchronous OpenRouterClient (used for aggregation).
-    :param async_client: An instance of an asynchronous OpenRouterClient.
+    :param provider: An instance of a OpenRouterProvider (used for aggregation).
+    :param async_provider: An instance of an AsyncOpenRouterProvider.
     :param consensus_config: An instance of ConsensusConfig.
     """
     response_data = {}
     response_data["initial_conversation"] = consensus_config.initial_prompt
 
     # Step 1: Initial round.
-    responses = await consensus.send_round(client, consensus_config)
-    aggregated_response = await aggregator.async_centralized_llm_aggregator(
-        client, consensus_config.aggregator_config, responses
+    responses = await send_round(provider, consensus_config)
+    aggregated_response = await async_centralized_llm_aggregator(
+        provider, consensus_config.aggregator_config, responses
     )
     logger.info(
         "initial response aggregation complete", aggregated_response=aggregated_response
@@ -42,11 +42,9 @@ async def run_consensus(
 
     # Step 2: Improvement rounds.
     for i in range(consensus_config.iterations):
-        responses = await consensus.send_round(
-            client, consensus_config, aggregated_response
-        )
-        aggregated_response = await aggregator.async_centralized_llm_aggregator(
-            client, consensus_config.aggregator_config, responses
+        responses = await send_round(provider, consensus_config, aggregated_response)
+        aggregated_response = await async_centralized_llm_aggregator(
+            provider, consensus_config.aggregator_config, responses
         )
         logger.info(
             "responses aggregated",
@@ -59,28 +57,28 @@ async def run_consensus(
 
     # Step 3: Save final consensus.
     output_file = config.data_path / "final_consensus.json"
-    saver.save_json(
+    save_json(
         response_data,
         output_file,
     )
     logger.info("saved consensus", output_file=output_file)
 
-    # Close the async client to release resources.
-    await client.close()
+    # Close the async provider to release resources.
+    await provider.close()
 
 
 def main() -> None:
     # Load the consensus configuration from input.json
-    config_json = loader.load_json(config.input_path / "input.json")
+    config_json = load_json(config.input_path / "input.json")
     consensus_config = ConsensusConfig.load_parameters(config_json)
 
-    # Initialize the OpenRouter client.
-    client = AsyncOpenRouterClient(
+    # Initialize the OpenRouter provider.
+    provider = AsyncOpenRouterProvider(
         api_key=config.open_router_api_key, base_url=config.open_router_base_url
     )
 
     # Run the consensus learning process with synchronous requests.
-    asyncio.run(run_consensus(client, consensus_config))
+    asyncio.run(run_consensus(provider, consensus_config))
 
 
 if __name__ == "__main__":

src/flare_ai_consensus/router/__init__.py

@@ -0,0 +1,10 @@
+from .base_router import ChatRequest, CompletionRequest, Message
+from .openrouter import AsyncOpenRouterProvider, OpenRouterProvider
+
+__all__ = [
+    "AsyncOpenRouterProvider",
+    "ChatRequest",
+    "CompletionRequest",
+    "Message",
+    "OpenRouterProvider",
+]