Added possibility to test connection

vianu/spock/__main__.py

@@ -2,19 +2,34 @@ from argparse import Namespace
 import asyncio
 from datetime import datetime
 import logging
+import os
 import sys
-from typing import List, Tuple
+from typing import List, Tuple, Dict, Any
+
+from dotenv import load_dotenv
 
 from vianu import LOG_FMT
-from vianu.spock.settings import SCRAPING_SOURCES, LOG_LEVEL
+from vianu.spock.settings import SCRAPING_SOURCES, LOG_LEVEL, MODEL_TEST_QUESTION
 from vianu.spock.src.cli import parse_args
 from vianu.spock.src.base import Setup, Document, SpoCK, FileHandler
 from vianu.spock.src import scraping as scp
 from vianu.spock.src import ner
 
-
 logging.basicConfig(format=LOG_FMT, level=LOG_LEVEL)
 logger = logging.getLogger(__name__)
+load_dotenv()
+
+
+def get_model_config() -> Dict[str, Dict[str, Any]]:
+    """Get model configuration."""
+    return {
+        'openai': {
+            'api_key': os.getenv('OPENAI_API_KEY'),
+        },
+        'llama': {
+            'base_url': os.getenv('OLLAMA_BASE_URL'),
+        },
+    }
 
 
 async def _orchestrator(
@@ -73,7 +88,7 @@ async def _orchestrator(
     await ner_queue.put(None)
 
 
-def setup_asyncio_framework(args_: Namespace) -> Tuple[asyncio.Queue, List[asyncio.Task], List[asyncio.Task], asyncio.Task]:
+def setup_asyncio_framework(args_: Namespace, model_config: Dict[str, Any]) -> Tuple[asyncio.Queue, List[asyncio.Task], List[asyncio.Task], asyncio.Task]:
     """Set up the asyncio framework for the SpoCK application."""
     # Set up arguments
     if args_.source is None:
@@ -86,7 +101,7 @@ def setup_asyncio_framework(args_: Namespace) -> Tuple[asyncio.Queue, List[async
 
     # Start tasks
     scp_tasks = scp.create_tasks(args_=args_, queue_in=src_queue, queue_out=scp_queue)
-    ner_tasks = ner.create_tasks(args_=args_, queue_in=scp_queue, queue_out=ner_queue)
+    ner_tasks = ner.create_tasks(args_=args_, queue_in=scp_queue, queue_out=ner_queue, model_config=model_config)
     orc_task = asyncio.create_task(
         _orchestrator(
             args_=args_,
@@ -125,10 +140,22 @@ async def main(args_: Namespace | None = None, save: bool = True) -> None:
         args_= parse_args(sys.argv[1:])
 
     logging.basicConfig(level=args_.log_level.upper(), format=LOG_FMT)
-    logger.info(f'starting SpoCK (args_={args_})')    
+    logger.info(f'starting SpoCK (args_={args_})')
+
+    # Test availability of NER model
+    model = args_.model
+    model_config = get_model_config()
+    try:
+        _ner = ner.NERFactory.create(model=model, config=model_config)
+        test_task = asyncio.create_task(_ner.test_model_endpoint())
+        test_answer = await test_task
+        logger.debug(f"test model endpoint of model='{model}': '{MODEL_TEST_QUESTION}' was answered with '{test_answer}'")
+    except Exception as e:
+        logger.error(f'could not reach model endpoint: {e}')
+        raise e
 
     # Set up async structure (scraping queue/tasks, NER queue/tasks, orchestrator task)
-    ner_queue, _, _, _ = setup_asyncio_framework(args_)
+    ner_queue, _, _, _ = setup_asyncio_framework(args_=args_, model_config=model_config)
 
     # Set up collector task and wait for it to finish
     # NOTE: if collector task is finished, the orchestrator is also finished (because of the sentinel in `ner_queue`)

vianu/spock/app/app.py

@@ -10,26 +10,30 @@ from dotenv import load_dotenv
 import gradio as gr
 
 from vianu.spock.settings import LOG_LEVEL, N_SCP_TASKS, N_NER_TASKS
-from vianu.spock.settings import LARGE_LANGUAGE_MODELS, SCRAPING_SOURCES, MAX_DOCS_SRC
+from vianu.spock.settings import LARGE_LANGUAGE_MODELS, SCRAPING_SOURCES, MAX_DOCS_SRC, MODEL_TEST_QUESTION
 from vianu.spock.settings import GRADIO_APP_NAME, GRADIO_SERVER_PORT, GRADIO_MAX_JOBS, GRADIO_UPDATE_INTERVAL
-from vianu.spock.settings import OLLAMA_BASE_URL_ENV_NAME, OPENAI_API_KEY_ENV_NAME
 from vianu.spock.src.base import Setup, SpoCK, SpoCKList, QueueItem         # noqa: F401
 from vianu import BaseApp
-from vianu.spock.__main__ import setup_asyncio_framework
+from vianu.spock.__main__ import get_model_config, setup_asyncio_framework
 import vianu.spock.app.formatter as fmt
+from vianu.spock.src.ner import NERFactory
 
 logger = logging.getLogger(__name__)
 load_dotenv()
 
 # App settings
 _ASSETS_PATH = Path(__file__).parents[1] / "assets"
-_UI_SETTINGS_LLM_CHOICES = [(name, value) for name, value in zip(['Ollama', 'OpenAI'], LARGE_LANGUAGE_MODELS)]
-_UI_SETTINGS_SOURCE_CHOICES = [(name, value) for name, value in zip(['PubMed', 'EMA', 'MHRA', 'FDA'], SCRAPING_SOURCES)]
+
+_UI_SETTINGS_LLM_CHOICES = [(name, value) for name, value in zip(['OpenAI', 'Ollama'], LARGE_LANGUAGE_MODELS)]
 if not len(_UI_SETTINGS_LLM_CHOICES) == len(LARGE_LANGUAGE_MODELS):
     raise ValueError('LARGE_LANGUAGE_MODELS and _UI_SETTINGS_LLM_CHOICES must have the same length')
+
+_UI_SETTINGS_SOURCE_CHOICES = [(name, value) for name, value in zip(['PubMed', 'EMA', 'MHRA', 'FDA'], SCRAPING_SOURCES)]
 if not len(_UI_SETTINGS_SOURCE_CHOICES) == len(SCRAPING_SOURCES):
     raise ValueError('SCRAPING_SOURCES and _UI_SETTINGS_SOURCE_CHOICES must have the same length')
 
+_MODEL_CONFIG = get_model_config()
+
 
 @dataclass
 class LocalState:
@@ -55,6 +59,12 @@ class SessionState:
     # Data
     is_running: bool = False
     spocks: SpoCKList = field(default_factory=list)
+
+    # Model config
+    model_config: Dict[str, Dict[str, Any]] = field(default_factory=lambda: _MODEL_CONFIG)
+    connection_is_valid: bool = False
+
+    # Indexes
     _index_running_spock: int | None = None
     _index_active_spock: int | None = None
 
@@ -113,25 +123,11 @@ class App(BaseApp):
         with gr.Column(scale=1):
             with gr.Accordion(label='LLM Endpoint'):
                 self._components['settings.llm_radio'] = gr.Radio(
-                    label='Model', show_label=False, choices=_UI_SETTINGS_LLM_CHOICES, value='llama', interactive=True
+                    label='Model', show_label=False, choices=_UI_SETTINGS_LLM_CHOICES, value='openai', interactive=True
                 )
 
-                # 'llama' specific settings
-                with gr.Group(visible=True) as self._components['settings.ollama_group']:
-                    value = os.environ.get("OLLAMA_BASE_URL")
-                    placeholder = 'base_url of ollama endpoint' if value is None else None
-                    gr.Markdown('---')
-                    self._components['settings.ollama_base_url'] = gr.Textbox(
-                        label='base_url',
-                        show_label=False,
-                        info='base_url',
-                        placeholder=placeholder,
-                        value=value,
-                        interactive=True,
-                    )
-
                 # 'openai' specific settings
-                with gr.Group(visible=False) as self._components['settings.openai_group']:
+                with gr.Group(visible=True) as self._components['settings.openai_group']:
                     value = os.environ.get("OPENAI_API_KEY")
                     placeholder = 'api_key of openai endpoint' if value is None else None
                     logger.debug(f'openai api_key={value}')
@@ -146,6 +142,23 @@ class App(BaseApp):
                         type='password',
                     )
 
+                # 'llama' specific settings
+                with gr.Group(visible=False) as self._components['settings.ollama_group']:
+                    value = os.environ.get("OLLAMA_BASE_URL")
+                    placeholder = 'base_url of ollama endpoint' if value is None else None
+                    gr.Markdown('---')
+                    self._components['settings.ollama_base_url'] = gr.Textbox(
+                        label='base_url',
+                        show_label=False,
+                        info='base_url',
+                        placeholder=placeholder,
+                        value=value,
+                        interactive=True,
+                    )
+                
+                self._components['settings.test_connection_button'] = gr.Button(value='Test connection', interactive=True)
+                
+
             with gr.Accordion(label='Sources', open=True):
                 self._components['settings.source'] = gr.CheckboxGroup(
                     label='Sources', show_label=False, choices=_UI_SETTINGS_SOURCE_CHOICES, value=SCRAPING_SOURCES, interactive=True
@@ -193,27 +206,48 @@ class App(BaseApp):
     # Helpers
     # --------------------------------------------------------------------------
     @staticmethod
-    def _show_llm_settings(llm: str) -> Tuple[dict[str, Any], dict[str, Any]]:
+    def _show_llm_settings(llm: str, session_state: SessionState) -> Tuple[dict[str, Any], dict[str, Any], SessionState]:
+        """Show the settings for the selected LLM model."""
         logger.debug(f'show {llm} model settings')
+        session_state.connection_is_valid = False
         if llm == 'llama':
-            return gr.update(visible=True), gr.update(visible=False)
+            return gr.update(visible=True), gr.update(visible=False), session_state
         elif llm == 'openai':
-            return gr.update(visible=False), gr.update(visible=True)
+            return gr.update(visible=False), gr.update(visible=True), session_state
         else:
-            return gr.update(visible=False), gr.update(visible=False)
+            return gr.update(visible=False), gr.update(visible=False), session_state
     
     @staticmethod
-    def _set_ollama_base_url(base_url: str) -> None:
-        """Setup ollama base_url as environment variable."""
-        logger.debug(f'set ollama base_url environment variable ({OLLAMA_BASE_URL_ENV_NAME}={base_url})')
-        os.environ[OLLAMA_BASE_URL_ENV_NAME] = base_url
+    def _set_openai_api_key(api_key: str, session_state: SessionState) -> SessionState:
+        """Setup openai api_key"""
+        log_key = '*****' if api_key else 'None'
+        logger.debug(f'set openai api_key {log_key}')
+        session_state.model_config['openai'] = {'api_key': api_key}
+        session_state.connection_is_valid = False
+        return session_state
     
     @staticmethod
-    def _set_openai_api_key(api_key: str) -> None:
-        """Setup openai api_key as environment variable."""
-        log_key = '*****' if api_key else 'None'
-        logger.debug(f'set openai api key (api_key={log_key})')
-        os.environ[OPENAI_API_KEY_ENV_NAME] = api_key
+    def _set_ollama_base_url(base_url: str, session_state: SessionState) -> SessionState:
+        """Setup ollama base_url"""
+        logger.debug(f'set ollama base_url={base_url}')
+        session_state.model_config['llama'] = {'base_url': base_url}
+        session_state.connection_is_valid = False
+        return session_state
+
+    @staticmethod
+    async def _test_connection(llm: str, session_state: SessionState) -> SessionState:
+        logger.debug(f'test connection to model={llm}')
+        try:
+            ner = NERFactory.create(model=llm, config=session_state.model_config)
+            test_task = asyncio.create_task(ner.test_model_endpoint())
+            test_answer = await test_task
+            gr.Info(f"connection to model={llm} is valid: '{MODEL_TEST_QUESTION}' was answered with '{test_answer}'")
+            session_state.connection_is_valid = True
+        except Exception as e:
+            session_state.connection_is_valid = False
+            raise gr.Error(f'connection to model={llm} failed: {e}')
+        return session_state
+
     
     @staticmethod
     def _feed_cards_to_ui(local_state: dict, session_state: SessionState) -> List[dict[str, Any]]:
@@ -243,15 +277,20 @@ class App(BaseApp):
         logger.debug(f'feeding details to UI (len(data)={len(active_spock.data)})')
         return fmt.get_details_html(active_spock.data)
 
-    @staticmethod
-    def _check_llm_settings(llm: str) -> None:
-        """Check if the LLM settings are set."""
-        if llm == 'llama':
-            if os.environ.get(OLLAMA_BASE_URL_ENV_NAME) is None:
-                raise gr.Error('Ollama base_url is not set (submit value with Enter)')
-        elif llm == 'openai':
-            if os.environ.get(OPENAI_API_KEY_ENV_NAME) is None:
-                raise gr.Error('OpenAI api_key is not set (submit value with Enter)')
+    async def _check_llm_settings(self, llm: str, session_state: SessionState) -> SessionState:
+        """Check if the LLM settings are set correcly."""
+        # Check if the settings are set
+        if llm == 'openai':
+            if session_state.model_config.get(llm) is None or session_state.model_config[llm].get('api_key') is None:
+                raise gr.Error('OpenAI api_key is not set')
+        elif llm == 'llama':
+            if session_state.model_config.get(llm) is None or session_state.model_config[llm].get('base_url') is None:
+                raise gr.Error('Ollama base_url is not set')
+            
+        # Check connection to the model
+        if not session_state.connection_is_valid:
+            session_state = await self._test_connection(llm=llm, session_state=session_state)
+        return session_state
 
     @staticmethod
     def _toggle_button(session_state: SessionState) -> Tuple[SessionState, dict[str, Any], dict[str, Any], dict[str, Any]]:
@@ -348,7 +387,7 @@ class App(BaseApp):
 
         # Setup asyncio tasks as in `vianu.spock.__main__`
         args_ = session_state.get_running_spock().setup.to_namespace()
-        ner_queue, scp_tasks, ner_tasks, orc_task = setup_asyncio_framework(args_=args_)
+        ner_queue, scp_tasks, ner_tasks, orc_task = setup_asyncio_framework(args_=args_, model_config=session_state.model_config)
         session_state.ner_queue = ner_queue
         session_state.scp_tasks = scp_tasks
         session_state.ner_tasks = ner_tasks
@@ -442,25 +481,42 @@ class App(BaseApp):
         """Choose LLM model show the correspoding settings."""
         self._components['settings.llm_radio'].change(
             fn=self._show_llm_settings,
-            inputs=self._components['settings.llm_radio'],
+            inputs=[
+                self._components['settings.llm_radio'],
+                self._session_state,
+            ],
             outputs=[
                 self._components['settings.ollama_group'],
                 self._components['settings.openai_group'],
+                self._session_state,
             ],
         )
+    
+    def _event_settings_openai(self):
+        """Callback of the openai settings."""
+        self._components['settings.openai_api_key'].change(
+            fn=self._set_openai_api_key,
+            inputs=[self._components['settings.openai_api_key'], self._session_state],
+            outputs=self._session_state,
+        )
 
     def _event_settings_ollama(self):
         """Callback of the ollama settings."""
-        self._components['settings.ollama_base_url'].submit(
+        self._components['settings.ollama_base_url'].change(
             fn=self._set_ollama_base_url,
-            inputs=self._components['settings.ollama_base_url'],
+            inputs=[self._components['settings.ollama_base_url'], self._session_state],
+            outputs=self._session_state,
         )
     
-    def _event_settings_openai(self):
-        """Callback of the openai settings."""
-        self._components['settings.openai_api_key'].submit(
-            fn=self._set_openai_api_key,
-            inputs=self._components['settings.openai_api_key'],
+    def _event_test_connection(self):
+        """Test the connection to the LLM model."""
+        self._components['settings.test_connection_button'].click(
+            fn=self._test_connection,
+            inputs=[
+                self._components['settings.llm_radio'],
+                self._session_state,
+            ],
+            outputs=self._session_state,
         )
 
     def _event_start_spock(self) -> None:
@@ -471,7 +527,11 @@ class App(BaseApp):
         gr.on(
             triggers=[search_term.submit, start_button.click],
             fn=self._check_llm_settings,
-            inputs=self._components['settings.llm_radio'],
+            inputs=[
+                self._components['settings.llm_radio'],
+                self._session_state,
+            ],
+            outputs=self._session_state,
         ).success(
             fn=self._toggle_button,
             inputs=self._session_state,
@@ -568,6 +628,7 @@ class App(BaseApp):
         self._event_choose_llm()
         self._event_settings_ollama()
         self._event_settings_openai()
+        self._event_test_connection()
 
         # Start/Stop events
         self._event_start_spock()

vianu/spock/app/formatter.py

@@ -11,6 +11,7 @@ JOBS_CONTAINER_CARD_TEMPLATE = """
 <div class="card" onclick="cardClickHandler(this)">
   <div class="title">{title} {status}</div>
   <div class="info">Date: {date}</div>
+  <div class="info">Model: {model}</div>
   <div class="info">Sources: {sources}</div>
   <div class="info">#docs: {n_doc} | #adr: {n_adr}</div>
 </div>
@@ -78,13 +79,14 @@ def _get_status_html(status: str) -> str:
 
 def get_job_card_html(card_nmbr: int, spock: SpoCK):
     """Get the HTML for the job card."""
-    job = spock.setup
+    setup = spock.setup
     data = spock.data
 
     title = spock.setup.term
     status = _get_status_html(spock.status)
-    sources = ", ".join(job.source)
-    date = job.submission.strftime(DATE_FORMAT)
+    date = setup.submission.strftime(DATE_FORMAT)
+    model = setup.model
+    sources = ", ".join(setup.source)
     n_doc = len(data)
     n_adr = sum([len(d.adverse_reactions) for d in data])
     return JOBS_CONTAINER_CARD_TEMPLATE.format(
@@ -92,6 +94,7 @@ def get_job_card_html(card_nmbr: int, spock: SpoCK):
         title=title,
         status=status,
         date=date,
+        model=model,
         sources=sources,
         n_doc=n_doc,
         n_adr=n_adr,

vianu/spock/settings.py

@@ -23,10 +23,11 @@ PUBMED_DB = 'pubmed'
 PUBMED_BATCH_SIZE = 20
 
 # NER settings
-LARGE_LANGUAGE_MODELS = ['llama', 'openai']
+LARGE_LANGUAGE_MODELS = ['openai', 'llama']
 MAX_TOKENS = 128.000
 LLAMA_MODEL='llama3.2'
 OPENAI_MODEL='gpt-4o'
 N_NER_TASKS = 2
 OLLAMA_BASE_URL_ENV_NAME = "OLLAMA_BASE_URL"
 OPENAI_API_KEY_ENV_NAME = "OPENAI_API_KEY"
+MODEL_TEST_QUESTION = 'Are you available?'

vianu/spock/src/cli.py

@@ -27,7 +27,7 @@ def parse_args(args_: Sequence) -> argparse.Namespace:
 
     # Add NER group
     ner_gp = parser.add_argument_group('ner')
-    ner_gp.add_argument('--model', '-m', type=str, choices=LARGE_LANGUAGE_MODELS, default='llama', help='NER model')
+    ner_gp.add_argument('--model', '-m', type=str, choices=LARGE_LANGUAGE_MODELS, default='openai', help='NER model')
     ner_gp.add_argument('--n-ner-tasks', metavar='', type=int, default=N_NER_TASKS, help='number of async ner tasks')
 
     return parser.parse_args(args_)

vianu/spock/src/ner.py

@@ -7,14 +7,12 @@ import os
 import re
 from typing import List
 
-from dotenv import load_dotenv
 from openai import AsyncOpenAI
 
 from vianu.spock.src.base import NamedEntity, QueueItem     # noqa: F401
-from vianu.spock.settings import N_CHAR_DOC_ID, LLAMA_MODEL, OPENAI_MODEL
+from vianu.spock.settings import N_CHAR_DOC_ID, LLAMA_MODEL, OPENAI_MODEL, MODEL_TEST_QUESTION
 
 logger = logging.getLogger(__name__)
-load_dotenv()
 
 NAMED_ENTITY_PROMPT = """
 You are an expert in Natural Language Processing. Your task is to identify named entities (NER) in a given text.
@@ -89,12 +87,27 @@ class NER(ABC):
                     "content": text,
                 },
             ]
+    
+    @staticmethod
+    def _get_test_messages(text: str) -> List[dict]:
+        return [
+                {
+                    "role": "system",
+                    "content": "Answer very briefly."},
+                {
+                    "role": "user",
+                    "content": text,
+                },
+            ]
 
 
     @abstractmethod
-    async def _get_ner_model_answer(self, text: str) -> str:
+    async def _get_model_answer(self, text: str) -> str:
         raise NotImplementedError('OpenAINER._get_ner_model_answer is not implemented yet')
 
+    @abstractmethod
+    async def test_model_endpoint(self) -> str:
+        pass
 
     async def apply(self, queue_in: asyncio.Queue, queue_out: asyncio.Queue) -> None:
         """Apply NER to a text received from input queue and put the results in an output queue."""
@@ -114,7 +127,7 @@ class NER(ABC):
             self.logger.debug(f'starting ner for item.id_={id_} (doc.id_={doc.id_[:N_CHAR_DOC_ID]})')
             try:
                 text = doc.text
-                content = await self._get_ner_model_answer(text=text)
+                content = await self._get_model_answer(text=text)
             except Exception as e:
                 self.logger.error(f'error during ner for item.id_={item.id_} (doc.id_={doc.id_[:N_CHAR_DOC_ID]}): {e}')
                 queue_in.task_done()
@@ -147,16 +160,38 @@ class NER(ABC):
             # Put the document in the output queue
             await queue_out.put(item)
             queue_in.task_done()
-            self.logger.info(f'finished NER task for item.id_={id_} (doc.id_={doc.id_[:N_CHAR_DOC_ID]})')
+            self.logger.info(f'finished NER task for item.id_={id_} (doc.id_={doc.id_[:N_CHAR_DOC_ID]})')  
 
 
-class OllamaNER(NER):
+class OpenAINER(NER):
 
-    def __init__(self, base_url: str, model: str):
+    def __init__(self, model: str, api_key: str):
         super().__init__()
-        self._base_url = base_url
         self._model = model
+        self._client = AsyncOpenAI(api_key=api_key)
+    
+    async def _chat_completion(self, messages: List[dict]) -> str:
+        chat_completion = await self._client.chat.completions.create(
+            messages=messages,
+            model=self._model,
+        )
+        return chat_completion.choices[0].message.content
 
+    async def _get_model_answer(self, text: str) -> str:
+        messages = self._get_messages(text=text)
+        return await self._chat_completion(messages=messages)
+    
+    async def test_model_endpoint(self) -> str:
+        messages = self._get_test_messages(text=MODEL_TEST_QUESTION)
+        return await self._chat_completion(messages=messages)
+
+
+class OllamaNER(NER):
+
+    def __init__(self, model: str, base_url: str):
+        super().__init__()
+        self._model = model
+        self._base_url = base_url
 
     def _get_http_data(self, text: str, stream: bool = False) -> dict:
         messages = self._get_messages(text=text)
@@ -166,9 +201,8 @@ class OllamaNER(NER):
             "stream": stream,
         }
         return data
-    
 
-    async def _get_ner_model_answer(self, text: str) -> str:
+    async def _get_model_answer(self, text: str) -> str:
         data = self._get_http_data(text=text)
         async with aiohttp.ClientSession() as session:
             url = f'{self._base_url}/api/chat/'
@@ -177,44 +211,57 @@ class OllamaNER(NER):
                 resp_json = await response.json()
                 content = resp_json['message']['content']
         return content
+    
+    async def test_model_endpoint(self) -> str:
+        messages = self._get_test_messages(text=MODEL_TEST_QUESTION)
+        data = {
+            'model': self._model,
+            "messages": messages,
+            'stream': False,
+        }
+        async with aiohttp.ClientSession() as session:
+            url = f'{self._base_url}/api/chat/'
+            async with session.post(url, json=data) as response:
+                response.raise_for_status()
+                resp_json = await response.json()
+                content = resp_json['message']['content']
+        return content
 
 
-class OpenAINER(NER):
+class NERFactory:
+    """Factory for NER models."""
 
-    def __init__(self, api_key: str, model: str):
-        super().__init__()
-        self._model = model
-        self._client = AsyncOpenAI(api_key=api_key)
+    @staticmethod
+    def create(model: str, config: dict) -> NER:
+        """Create a NER model from the `model` keyword."""
+
+        # create a :class:`OllamaNER` instance
+        if model == 'openai':
+            api_key = config[model].get('api_key')
+            if api_key is None:
+                api_key = os.environ.get('OPENAI_API_KEY')
+            if api_key is None:
+                raise ValueError("The api_key for the OpenAI client is missing (set it by the OPENAI_API_KEY environment variable or in the settings)")
+            ner = OpenAINER(model=OPENAI_MODEL, api_key=api_key)
+        
+        elif model == 'llama':
+            base_url = config[model].get('base_url')
+            if base_url is None:
+                base_url = os.environ.get('OLLAMA_BASE_URL')
+            if base_url is None:
+                raise ValueError("The base_url for the ollama endpoint is missing")
+            ner = OllamaNER(model=LLAMA_MODEL, base_url=base_url)
+        else:
+            raise ValueError(f"unknown ner model '{model}'")
     
-
-    async def _get_ner_model_answer(self, text: str) -> str:
-        messages = self._get_messages(text=text)
-        chat_completion = await self._client.chat.completions.create(
-            messages=messages,
-            model=self._model,
-        )
-        return chat_completion.choices[0].message.content
+        return ner
 
 
-def create_tasks(args_: Namespace, queue_in: asyncio.Queue, queue_out: asyncio.Queue) -> List[asyncio.Task]:
+def create_tasks(args_: Namespace, queue_in: asyncio.Queue, queue_out: asyncio.Queue, model_config: dict) -> List[asyncio.Task]:
     """Create asyncio NER tasks."""
     n_ner_tasks = args_.n_ner_tasks
     model = args_.model
-    
-    if model == 'llama':
-        base_url = os.environ.get('OLLAMA_BASE_URL')
-        if base_url is None:
-            raise EnvironmentError("The ollama endpoint must be set by the OLLAMA_ENDPOINT environment variable")
-        ner = OllamaNER(base_url=base_url, model=LLAMA_MODEL)
-
-    elif model == 'openai':
-        api_key = os.environ.get('OPENAI_API_KEY')
-        if api_key is None:
-            raise EnvironmentError("The api_key for the OpenAI client must be set by the OPENAI_API_KEY environment variable")
-        ner = OpenAINER(api_key=api_key, model=OPENAI_MODEL)
-
-    else:
-        raise ValueError(f"unknown ner model '{args_.model}'")
+    ner = NERFactory.create(model=model, config=model_config)
     
     logger.info(f'setting up {n_ner_tasks} NER task(s)')
     tasks = [asyncio.create_task(ner.apply(queue_in=queue_in, queue_out=queue_out)) for _ in range(n_ner_tasks)]

vianu/spock/src/scraping.py

@@ -9,7 +9,7 @@ The module contains three main classes:
 from abc import ABC, abstractmethod
 from argparse import Namespace
 import asyncio
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime
 from io import BytesIO
 import logging
@@ -284,9 +284,9 @@ class PubmedScraper(Scraper):
 @dataclass
 class SearchResults:
     """Class for storing the search results from different databases."""
-    count: int | None
-    n_pages: int | None
-    items: List[Tag]
+    count: int | None = None
+    n_pages: int | None = None
+    items: List[Tag] = field(default_factory=list)
 
 
 class EMAScraper(Scraper):