tried to fix eval

Author: xuayan-nokia

wren-ai-service/eval/pipelines.py

@@ -2,14 +2,35 @@
 import os
 import re
 import sys
+import uuid
 from abc import abstractmethod
 from pathlib import Path
 from typing import Any, Dict, List, Literal
 
 import orjson
+import json
 from haystack import Document
 from langfuse.decorators import langfuse_context, observe
 from tqdm.asyncio import tqdm_asyncio
+from src.config import settings
+from src.providers import generate_components
+from src.web.v1.services.semantics_preparation import (
+    SemanticsPreparationRequest,
+    SemanticsPreparationService,
+)
+from src.web.v1.services.ask import (
+    AskRequest,
+    AskResultRequest,
+    AskResultResponse,
+    AskService,
+)
+from src.pipelines.generation import (
+    data_assistance,
+    intent_classification,
+    sql_correction,
+    sql_generation,
+)
+from src.pipelines.retrieval import historical_question, retrieval
 
 sys.path.append(f"{Path().parent.resolve()}")
 
@@ -32,15 +53,15 @@
 from src.core.engine import Engine
 from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider
 from src.pipelines.generation import sql_generation
-from src.pipelines.indexing import indexing
 from src.pipelines.retrieval import retrieval
+from src.pipelines import indexing
 
 
-def deploy_model(mdl: str, pipe: indexing.Indexing) -> None:
-    async def wrapper():
-        await pipe.run(orjson.dumps(mdl).decode())
+# def deploy_model(mdl: str, pipe: indexing.Indexing) -> None:
+#     async def wrapper():
+#         await pipe.run(orjson.dumps(mdl).decode())
 
-    asyncio.run(wrapper())
+#     asyncio.run(wrapper())
 
 
 def extract_units(docs: list) -> list:
@@ -107,6 +128,7 @@ def split(queries: list, batch_size: int) -> list[list]:
             ]
 
         async def wrapper(batch: list):
+            # self() will call sub-class's __call__ in every service
             tasks = [self(query) for query in batch]
             results = await tqdm_asyncio.gather(*tasks, desc="Generating Predictions")
             await asyncio.sleep(self._batch_interval)
@@ -188,7 +210,7 @@ def __init__(
             embedder_provider=embedder_provider,
             document_store_provider=document_store_provider,
         )
-        deploy_model(mdl, _indexing)
+        # deploy_model(mdl, _indexing)
 
         self._retrieval = retrieval.Retrieval(
             llm_provider=llm_provider,
@@ -288,36 +310,82 @@ def mertics(
         }
 
 
+
 class AskPipeline(Eval):
+    def indexing_service(self):
+
+        return SemanticsPreparationService(
+            {
+                "db_schema": indexing.DBSchema(
+                    **self.pipe_components["db_schema_indexing"],
+                ),
+                "historical_question": indexing.HistoricalQuestion(
+                    **self.pipe_components["historical_question_indexing"],
+                ),
+                "table_description": indexing.TableDescription(
+                    **self.pipe_components["table_description_indexing"],
+                ),
+            }
+        )
+
+    def ask_service(self):
+
+        return AskService(
+            {
+                "intent_classification": intent_classification.IntentClassification(
+                    **self.pipe_components["intent_classification"],
+                ),
+                "data_assistance": data_assistance.DataAssistance(
+                    **self.pipe_components["data_assistance"],
+                ),
+                "retrieval": retrieval.Retrieval(
+                    **self.pipe_components["db_schema_retrieval"],
+                ),
+                "historical_question": historical_question.HistoricalQuestion(
+                    **self.pipe_components["historical_question_retrieval"],
+                ),
+                "sql_generation": sql_generation.SQLGeneration(
+                    **self.pipe_components["sql_generation"],
+                ),
+                "sql_correction": sql_correction.SQLCorrection(
+                    **self.pipe_components["sql_correction"],
+                ),
+            }
+        )
+    def dict_to_string(self, d: dict) -> str:
+        if not isinstance(d, dict):
+            return str(d)
+
+        result = "{"
+        for key, value in d.items():
+            result += f"'{key}': {self.dict_to_string(value)}, "
+        result = result.rstrip(", ") + "}"
+        return result
+
     def __init__(
         self,
         meta: dict,
         mdl: dict,
-        llm_provider: LLMProvider,
-        embedder_provider: EmbedderProvider,
-        document_store_provider: DocumentStoreProvider,
-        engine: Engine,
-        **kwargs,
+        service_metadata,
+        pipe_components,
     ):
         super().__init__(meta, 3)
-
-        document_store_provider.get_store(recreate_index=True)
-        _indexing = indexing.Indexing(
-            embedder_provider=embedder_provider,
-            document_store_provider=document_store_provider,
-        )
-        deploy_model(mdl, _indexing)
-
+        self.service_metadata = service_metadata
+
+        # document_store_provider.get_store(recreate_index=True)
+        # _indexing = indexing.Indexing(
+        #     embedder_provider=embedder_provider,
+        #     document_store_provider=document_store_provider,
+        # )
+        # deploy_model(mdl, _indexing)
+        self.pipe_components = pipe_components
+        self.project_id = str(uuid.uuid4())
+        self.indexing_service_var = self.indexing_service()
+        self.mdl_str_var = json.dumps(mdl)
+        self.ask_service_var = self.ask_service()
+        self.service_metadata = service_metadata
         self._mdl = mdl
-        self._retrieval = retrieval.Retrieval(
-            llm_provider=llm_provider,
-            embedder_provider=embedder_provider,
-            document_store_provider=document_store_provider,
-        )
-        self._generation = sql_generation.SQLGeneration(
-            llm_provider=llm_provider,
-            engine=engine,
-        )
+        self.mdl_hash = str(hash(self.mdl_str_var))
 
     async def _flat(self, prediction: dict, actual: str) -> dict:
         prediction["actual_output"] = actual
@@ -327,17 +395,54 @@ async def _flat(self, prediction: dict, actual: str) -> dict:
         return prediction
 
     async def _process(self, prediction: dict, **_) -> dict:
-        result = await self._retrieval.run(query=prediction["input"])
-        documents = result.get("construct_retrieval_results", [])
-        actual_output = await self._generation.run(
+
+        await self.indexing_service_var.prepare_semantics(
+            SemanticsPreparationRequest(
+                mdl=self.mdl_str_var,
+                mdl_hash=self.mdl_hash,
+                project_id=self.project_id
+            ),
+            service_metadata=self.service_metadata,
+        )
+
+        # asking
+        ask_request = AskRequest(
             query=prediction["input"],
-            contexts=documents,
-            samples=prediction["samples"],
-            exclude=[],
+            mdl_hash=self.mdl_hash,
+            project_id = self.project_id,
+
+        )
+        ask_request.query_id = str(uuid.uuid4())
+        await self.ask_service_var.ask(ask_request, service_metadata=self.service_metadata)
+        # getting ask result
+        ask_result_response = self.ask_service_var.get_ask_result(
+            AskResultRequest(
+                query_id=ask_request.query_id,
+            )
         )
 
-        prediction["actual_output"] = actual_output
-        prediction["retrieval_context"] = extract_units(documents)
+        while (
+            ask_result_response.status != "finished"
+            and ask_result_response.status != "failed"
+        ):
+            # getting ask result
+            ask_result_response = self.ask_service_var.get_ask_result(
+                AskResultRequest(
+                    query_id=ask_request.query_id,
+                )
+            )
+
+        # result = await self._retrieval.run(query=prediction["input"])
+        # documents = result.get("construct_retrieval_results", [])
+        # actual_output = await self._generation.run(
+        #     query=prediction["input"],
+        #     contexts=documents,
+        #     samples=prediction["samples"],
+        #     exclude=[],
+        # )
+
+        prediction["actual_output"] = ask_result_response.response[0].sql
+        #prediction["retrieval_context"] = extract_units(documents)
 
         return prediction
 
@@ -377,9 +482,10 @@ def init(
     name: Literal["retrieval", "generation", "ask"],
     meta: dict,
     mdl: dict,
-    providers: Dict[str, Any],
+    service_metadata,
+    pipe_components: Dict[str, Any],
 ) -> Eval:
-    args = {"meta": meta, "mdl": mdl, **providers}
+    args = {"meta": meta, "mdl": mdl, "service_metadata":service_metadata,"pipe_components":pipe_components}
     match name:
         case "retrieval":
             return RetrievalPipeline(**args)

wren-ai-service/eval/prediction.py

@@ -14,13 +14,17 @@
 from tomlkit import document, dumps
 
 sys.path.append(f"{Path().parent.resolve()}")
+from src.config import settings
+from src.providers import generate_components
 import eval.pipelines as pipelines
-import src.providers as provider
 import src.utils as utils
 from eval.utils import parse_toml
 from src.core.engine import EngineConfig
 from src.core.provider import EmbedderProvider, LLMProvider
-
+from src.globals import (
+    create_service_container,
+    create_service_metadata,
+)
 
 def generate_meta(
     path: str,
@@ -46,10 +50,10 @@ def generate_meta(
         "commit": obtain_commit_hash(),
         "embedding_model": embedder_provider.get_model(),
         "generation_model": llm_provider.get_model(),
-        "column_indexing_batch_size": int(os.getenv("COLUMN_INDEXING_BATCH_SIZE"))
+        "column_indexing_batch_size": int(settings.column_indexing_batch_size)
         or 50,
-        "table_retrieval_size": int(os.getenv("TABLE_RETRIEVAL_SIZE")) or 10,
-        "table_column_retrieval_size": int(os.getenv("TABLE_COLUMN_RETRIEVAL_SIZE"))
+        "table_retrieval_size": int(settings.table_retrieval_size) or 10,
+        "table_column_retrieval_size": int(settings.table_column_retrieval_size)
         or 100,
         "pipeline": pipe,
         "batch_size": os.getenv("BATCH_SIZE") or 4,
@@ -84,11 +88,11 @@ def write_prediction(
 def obtain_commit_hash() -> str:
     repo = Repo(search_parent_directories=True)
 
-    if repo.untracked_files:
-        raise Exception("There are untracked files in the repository.")
+    # if repo.untracked_files:
+    #     raise Exception("There are untracked files in the repository.")
 
-    if repo.index.diff(None):
-        raise Exception("There are uncommitted changes in the repository.")
+    # if repo.index.diff(None):
+    #     raise Exception("There are uncommitted changes in the repository.")
 
     branch = repo.active_branch
     return f"{repo.head.commit}@{branch.name}"
@@ -138,12 +142,12 @@ def init_providers(mdl: dict) -> dict:
     if engine_config is None:
         raise ValueError("Invalid datasource")
 
-    providers = provider.init_providers(engine_config=engine_config)
+    providers_inner = provider.init_providers(engine_config=engine_config)
     return {
-        "llm_provider": providers[0],
-        "embedder_provider": providers[1],
-        "document_store_provider": providers[2],
-        "engine": providers[3],
+        "llm_provider": providers_inner[0],
+        "embedder_provider": providers_inner[1],
+        "document_store_provider": providers_inner[2],
+        "engine": providers_inner[3],
     }
 
 
@@ -174,23 +178,24 @@ def parse_args() -> Tuple[str]:
     utils.init_langfuse()
 
     dataset = parse_toml(path)
-    providers = init_providers(dataset["mdl"])
 
+    pipe_components = generate_components(settings.components)
     meta = generate_meta(
         path=path,
         dataset=dataset,
         pipe=pipe_name,
-        **providers,
+        **pipe_components["db_schema_retrieval"],
     )
-
+    service_metadata = create_service_metadata(pipe_components)
     pipe = pipelines.init(
         pipe_name,
         meta,
         mdl=dataset["mdl"],
-        providers=providers,
+        service_metadata=service_metadata,
+        pipe_components=pipe_components,
     )
 
-    predictions = pipe.predict(dataset["eval_dataset"])
+    predictions = pipe.predict([dataset["eval_dataset"][0]])
     meta["expected_batch_size"] = meta["query_count"] * pipe.candidate_size
     meta["actual_batch_size"] = len(predictions) - meta["query_count"]
 

wren-ai-service/src/globals.py

@@ -67,6 +67,13 @@ class ServiceContainer:
 class ServiceMetadata:
     pipes_metadata: dict
     service_version: str
+    def get(self, key: str):
+        if key=="service_version":
+            return self.service_version
+        elif key=="pipes_metadata":
+            return self.pipes_metadata
+        else:
+            return None
 
 
 def create_service_container(