andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Improve REST API

Browse source
This commit is contained in:
Andras Schmelczer 2022-06-03 19:29:48 +02:00
parent 014f84f390
commit b125ebc08c
31 changed files with 343 additions and 1678906 deletions
Download patch file Download diff file Expand all files Collapse all files

3
.vscode/settings.json vendored
View file

@ -3,6 +3,7 @@
"boto",
"botocore",
"fastapi",
"gridfs",
"iloc",
"inplace",
"ipynb",
@ -10,8 +11,10 @@
"matplotlib",
"nbconvert",
"plotly",
"proba",
"psutil",
"pydantic",
"pymongo",
"pyplot",
"redoc",
"sklearn",

82
examples/simple/__deploy__.py
View file

@ -1,82 +0,0 @@
#!/usr/bin/env python
# coding: utf-8
# # Train a domain classifier on the [semantic scholar dataset](https://api.semanticscholar.org/corpus)
#
# ## Part 3: Create production inference function
#
# In the [previous notebook](train.ipynb), we trained our AI model. Now, it's time to create **G**eneral **R**obust **E**nd-to-end **A**utomated **T**rustworthy deployment from it using the `GreatAI` Python package.
# In[1]:
import re
from typing import List
from great_ai import ClassificationOutput, GreatAI, use_model
from great_ai.utilities.clean import clean
from sklearn.pipeline import Pipeline
# In[2]:
@GreatAI.deploy
@use_model("small-domain-prediction-v2", version="latest")
def predict_domain(
text: str, model: Pipeline, target_confidence: int = 20
) -> List[ClassificationOutput]:
"""
Predict the scientific domain of the input text.
Return labels until their sum likelihood is larger than target_confidence.
"""
assert 0 <= target_confidence <= 100, "invalid argument"
preprocessed = re.sub(r"[^a-zA-Z ]", "", clean(text, convert_to_ascii=True))
features = model.named_steps["vectorizer"].transform([preprocessed])
prediction = model.named_steps["classifier"].predict_proba(features)[0]
best_classes = sorted(enumerate(prediction), key=lambda v: v[1], reverse=True)
results: List[ClassificationOutput] = []
for class_index, probability in best_classes:
results.append(
ClassificationOutput(
label=model.named_steps["classifier"].classes_[class_index],
confidence=round(probability * 100),
explanation=[
word
for _, word in sorted(
(
(weight, word)
for weight, word, count in zip(
model.named_steps["classifier"].feature_log_prob_[
class_index
],
model.named_steps["vectorizer"].get_feature_names_out(),
features.A[0],
)
if count > 0
),
reverse=True,
)
][:5],
)
)
if sum(r.confidence for r in results) >= target_confidence:
break
return results
# In[3]:
result = predict_domain(
"""
State-of-the-art methods for zero-shot visual recognition formulate learning as a joint embedding problem of images and side information. In these formulations the current best complement to visual features are attributes: manually encoded vectors describing shared characteristics among categories. Despite good performance, attributes have limitations: (1) finer-grained recognition requires commensurately more, and (2) attributes do not provide a natural language interface. We propose to overcome these limitations by training neural language models from scratch; i.e. without pre-training and only consuming words and characters. Our proposed models train end-to-end to align with the fine-grained and category-specific content of images. Natural language provides a flexible and compact way of encoding only the salient visual aspects for distinguishing categories. By training on raw text, our model can do inference on raw text as well, providing humans a familiar mode both for annotation and retrieval. Our model achieves strong performance on zero-shot text-based image retrieval and significantly outperforms the attribute-based state-of-the-art for zero-shot classification on the CaltechUCSD Birds 200-2011 dataset. """
)
from pprint import pprint
pprint(result.dict(), width=120)

30
examples/simple/deploy.ipynb
View file

@ -33,18 +33,18 @@
"name": "stderr",
"output_type": "stream",
"text": [
"\u001b[38;5;226m2022-05-28 15:02:20,852 | WARNING | Environment variable ENVIRONMENT is not set, defaulting to development mode ‼️\u001b[0m\n",
"\u001b[38;5;39m2022-05-28 15:02:20,853 | INFO | Options: configured ✅\u001b[0m\n",
"\u001b[38;5;39m2022-05-28 15:02:21,168 | INFO | Latest version of small-domain-prediction-v2 is 8 (from versions: 3, 4, 5, 6, 7, 8)\u001b[0m\n",
"\u001b[38;5;39m2022-05-28 15:02:21,169 | INFO | File small-domain-prediction-v2-8 found in cache\u001b[0m\n"
"\u001b[38;5;226m2022-05-28 18:17:13,344 | WARNING | Environment variable ENVIRONMENT is not set, defaulting to development mode ‼️\u001b[0m\n",
"\u001b[38;5;39m2022-05-28 18:17:13,346 | INFO | Options: configured ✅\u001b[0m\n",
"\u001b[38;5;39m2022-05-28 18:17:13,694 | INFO | Latest version of small-domain-prediction is 10 (from versions: 9, 10)\u001b[0m\n",
"\u001b[38;5;39m2022-05-28 18:17:13,694 | INFO | File small-domain-prediction-10 found in cache\u001b[0m\n"
]
}
],
"source": [
"@GreatAI.deploy\n",
"@use_model(\"small-domain-prediction-v2\", version=\"latest\")\n",
"@use_model(\"small-domain-prediction\", version=\"latest\")\n",
"def predict_domain(\n",
" text: str, model: Pipeline, target_confidence: int = 20\n",
" text: str, model: Pipeline, target_confidence: int = 50\n",
") -> List[ClassificationOutput]:\n",
" \"\"\"\n",
" Predict the scientific domain of the input text.\n",
@ -52,7 +52,7 @@
" \"\"\"\n",
" assert 0 <= target_confidence <= 100, \"invalid argument\"\n",
"\n",
" preprocessed = re.sub(r\"[^a-zA-Z ]\", \"\", clean(text, convert_to_ascii=True))\n",
" preprocessed = re.sub(r\"[^a-zA-Z\\s]\", \"\", clean(text, convert_to_ascii=True))\n",
" features = model.named_steps[\"vectorizer\"].transform([preprocessed])\n",
" prediction = model.named_steps[\"classifier\"].predict_proba(features)[0]\n",
"\n",
@ -99,12 +99,11 @@
"name": "stdout",
"output_type": "stream",
"text": [
"{'created': '2022-05-28T13:02:22.035886',\n",
" 'evaluation': None,\n",
" 'evaluation_id': 'c035d78e-dea9-415f-8f93-0dc4cd8dd7b5',\n",
"{'created': '2022-05-28T16:17:14.581693',\n",
" 'exception': None,\n",
" 'execution_time_ms': 94.814,\n",
" 'logged_values': {'arg:predict_domain:target_confidence': 20,\n",
" 'execution_time_ms': 93.638,\n",
" 'feedback': None,\n",
" 'logged_values': {'arg:predict_domain:target_confidence': 50,\n",
" 'arg:predict_domain:text': '\\n'\n",
" ' State-of-the-art methods for zero-shot visual recognition formulate '\n",
" 'learning as a joint embedding problem of images and side information. '\n",
@ -125,10 +124,11 @@
" 'outperforms the attribute-based state-of-the-art for zero-shot '\n",
" 'classification on the CaltechUCSD Birds 200-2011 dataset. ',\n",
" 'arg:predict_domain:text:length': 1236},\n",
" 'models': [{'key': 'small-domain-prediction-v2', 'version': 8}],\n",
" 'models': [{'key': 'small-domain-prediction', 'version': 10}],\n",
" 'output': [{'confidence': 99.0,\n",
" 'explanation': ['information', 'model', 'learning', 'proposed', 'image'],\n",
" 'label': 'Computer Science'}]}\n"
" 'label': 'Computer Science'}],\n",
" 'trace_id': 'd35fcb96-0a95-45f8-93e4-8967160b17dd'}\n"
]
}
],
@ -140,7 +140,7 @@
"\n",
"from pprint import pprint\n",
"\n",
"pprint(result.dict(), width=120)"
"# pprint(result.dict(), width=120)"
]
}
],

1678667
examples/simple/train.ipynb
View file

File diff suppressed because one or more lines are too long

10
great_ai/src/great_ai/great_ai/dashboard/create_dash_app.py
View file

@ -33,7 +33,7 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
],
)
documents = get_context().persistence.get_documents()
documents = get_context().tracing_database.query()
df = pd.DataFrame(documents)
execution_time_histogram = dcc.Graph(config={"displaylogo": False})
@ -95,11 +95,11 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
]
non_null_conjunctive_filters = [f for f in conjunctive_filters if f is not None]
return get_context().persistence.query(
conjunctive_filters=non_null_conjunctive_filters,
sort_by=sort_by,
return get_context().tracing_database.query(
skip=page_current * page_size,
take=page_size,
conjunctive_filters=non_null_conjunctive_filters,
sort_by=sort_by,
)
@app.callback(
@ -114,7 +114,7 @@ def create_dash_app(function_name: str, function_docs: str) -> Flask:
]
non_null_conjunctive_filters = [f for f in conjunctive_filters if f is not None]
rows = get_context().persistence.query(
rows = get_context().tracing_database.query(
conjunctive_filters=non_null_conjunctive_filters
)

145
great_ai/src/great_ai/great_ai/deploy/great_ai.py
View file

@ -1,27 +1,11 @@
import inspect
from functools import lru_cache, partial, wraps
from pathlib import Path
from typing import (
Any,
Callable,
Dict,
Iterable,
List,
Optional,
Sequence,
Type,
Union,
cast,
)
from typing import Any, Callable, Iterable, Optional, Sequence, Type, Union, cast
from fastapi import APIRouter, FastAPI, HTTPException, status
from fastapi.middleware.wsgi import WSGIMiddleware
from fastapi.openapi.docs import get_swagger_ui_html
from fastapi.responses import RedirectResponse
from fastapi.staticfiles import StaticFiles
from fastapi import APIRouter, FastAPI, status
from pydantic import BaseModel, create_model
from starlette.responses import HTMLResponse
from great_ai.great_ai.views.cache_statistics import CacheStatistics
from great_ai.utilities.parallel_map import parallel_map
from ..constants import METRICS_PATH
@ -34,36 +18,32 @@ from ..helper import (
use_http_exceptions,
)
from ..parameters import automatically_decorate_parameters
from ..tracing import TracingContext
from ..views import (
ApiMetadata,
EvaluationFeedbackRequest,
HealthCheckResponse,
Query,
Trace,
from ..tracing.tracing_context import TracingContext
from ..views import ApiMetadata, HealthCheckResponse, Trace
from .routes import (
bootstrap_docs_endpoints,
bootstrap_feedback_endpoints,
bootstrap_trace_endpoints,
)
PATH = Path(__file__).parent.resolve()
class GreatAI:
def __init__(self, func: Callable[..., Any]):
self._func = automatically_decorate_parameters(func)
self._func = freeze_arguments(
lru_cache(get_context().prediction_cache_size)(self._func)
)
get_function_metadata_store(self._func).is_finalised = True
wraps(func)(self)
self.app = FastAPI(
title=self.name,
version=self.version,
description=self.documentation,
description=self.documentation
+ f"\n\n Find out more on the [metrics page]({METRICS_PATH}).",
docs_url=None,
redoc_url=None,
)
@freeze_arguments
@lru_cache(get_context().prediction_cache_size)
def __call__(self, *args: Any, **kwargs: Any) -> Trace:
with TracingContext() as t:
result = self._func(*args, **kwargs)
@ -103,7 +83,7 @@ class GreatAI:
batch: Iterable[Any],
concurrency: Optional[int] = None,
) -> Sequence[Trace]:
if not get_context().persistence.is_threadsafe:
if not get_context().tracing_database.is_threadsafe:
concurrency = 1
get_context().logger.warning("Concurrency is ignored")
@ -132,14 +112,17 @@ class GreatAI:
def _bootstrap_rest_api(self, disable_docs: bool, disable_metrics: bool) -> None:
self._bootstrap_prediction_endpoints()
self._bootstrap_feedback_endpoints()
self._bootstrap_meta_endpoints()
if not disable_docs:
self._bootstrap_docs_endpoints()
bootstrap_docs_endpoints(self.app)
if not disable_metrics:
self._bootstrap_metrics_endpoints()
dash_app = create_dash_app(self._func.__name__, self.documentation)
bootstrap_trace_endpoints(self.app, dash_app)
bootstrap_feedback_endpoints(self.app)
self._bootstrap_meta_endpoints()
def _bootstrap_prediction_endpoints(self) -> None:
router = APIRouter(
@ -154,19 +137,6 @@ class GreatAI:
def predict(input_value: schema) -> Trace: # type: ignore
return self(**cast(BaseModel, input_value).dict())
@router.get(
"/:prediction_id", response_model=Trace, status_code=status.HTTP_200_OK
)
def get_prediction(prediction_id: str) -> Trace:
result = get_context().persistence.get_trace(prediction_id)
if result is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
return result
@router.delete("/:prediction_id", status_code=status.HTTP_204_NO_CONTENT)
def delete_prediction(prediction_id: str) -> None:
get_context().persistence.delete_trace(prediction_id)
self.app.include_router(router)
def _get_schema(self) -> Type[BaseModel]:
@ -183,26 +153,6 @@ class GreatAI:
schema: Type[BaseModel] = create_model("InputModel", **parameters) # type: ignore
return schema
def _bootstrap_feedback_endpoints(self) -> None:
router = APIRouter(
prefix="/predictions/:prediction_id/feedback",
tags=["feedback"],
)
@router.put("/", status_code=status.HTTP_202_ACCEPTED)
def set_feedback(prediction_id: str, input: EvaluationFeedbackRequest) -> None:
get_context().persistence.add_feedback(prediction_id, input.evaluation)
@router.get("/", status_code=status.HTTP_200_OK)
def get_feedback(prediction_id: str) -> Any:
return get_context().persistence.get_feedback(prediction_id)
@router.delete("/", status_code=status.HTTP_200_OK)
def delete_feedback(prediction_id: str) -> Any:
get_context().persistence.delete_feedback(prediction_id)
self.app.include_router(router)
def _bootstrap_meta_endpoints(self) -> None:
router = APIRouter(
tags=["meta"],
@ -210,53 +160,24 @@ class GreatAI:
@router.get("/health", status_code=status.HTTP_200_OK)
def check_health() -> HealthCheckResponse:
return HealthCheckResponse(is_healthy=True)
hits, misses, maxsize, cache_size = self.__call__.cache_info() # type: ignore
cache_statistics = CacheStatistics(
hits=hits, misses=misses, size=cache_size, max_size=maxsize
)
return HealthCheckResponse(
is_healthy=True, cache_statistics=cache_statistics
)
@router.get(
"/version", response_model=ApiMetadata, status_code=status.HTTP_200_OK
)
def get_version() -> ApiMetadata:
return ApiMetadata(
name=self.name, version=self.version, documentation=self.documentation
)
self.app.include_router(router)
def _bootstrap_docs_endpoints(self) -> None:
@self.app.get("/docs", include_in_schema=False)
def custom_swagger_ui_html() -> HTMLResponse:
return get_swagger_ui_html(openapi_url="openapi.json", title=self.app.title)
@self.app.get("/docs/index.html", include_in_schema=False)
def redirect_to_docs() -> RedirectResponse:
return RedirectResponse("/docs")
def _bootstrap_metrics_endpoints(self) -> None:
dash_app = create_dash_app(self._func.__name__, self.documentation)
self.app.mount(METRICS_PATH, WSGIMiddleware(dash_app))
@self.app.get("/", include_in_schema=False)
def redirect_to_entrypoint() -> RedirectResponse:
return RedirectResponse("/metrics")
self.app.mount(
"/assets",
StaticFiles(directory=PATH / "../dashboard/assets"),
name="static",
)
router = APIRouter(
prefix="/metrics",
tags=["metrics"],
)
@router.post("/query", status_code=status.HTTP_200_OK)
def query_metrics(query: Query) -> List[Dict[str, Any]]:
return get_context().persistence.query(
conjunctive_filters=query.filter,
sort_by=query.sort,
skip=query.skip,
take=query.take,
name=self.name,
version=self.version,
documentation=self.documentation,
configuration=get_context().to_flat_dict(),
)
self.app.include_router(router)

3
great_ai/src/great_ai/great_ai/deploy/routes/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .bootstrap_docs_endpoints import bootstrap_docs_endpoints
from .bootstrap_feedback_endpoints import bootstrap_feedback_endpoints
from .bootstrap_trace_endpoints import bootstrap_trace_endpoints

14
great_ai/src/great_ai/great_ai/deploy/routes/bootstrap_docs_endpoints.py Normal file
View file

@ -0,0 +1,14 @@
from fastapi import FastAPI
from fastapi.openapi.docs import get_swagger_ui_html
from fastapi.responses import RedirectResponse
from starlette.responses import HTMLResponse
def bootstrap_docs_endpoints(app: FastAPI) -> None:
@app.get("/docs", include_in_schema=False)
def custom_swagger_ui_html() -> HTMLResponse:
return get_swagger_ui_html(openapi_url="openapi.json", title=app.title)
@app.get("/docs/index.html", include_in_schema=False)
def redirect_to_docs() -> RedirectResponse:
return RedirectResponse("/docs")

40
great_ai/src/great_ai/great_ai/deploy/routes/bootstrap_feedback_endpoints.py Normal file
View file

@ -0,0 +1,40 @@
from typing import Any
from fastapi import APIRouter, FastAPI, HTTPException, Response, status
from ...context import get_context
from ...views import EvaluationFeedbackRequest
def bootstrap_feedback_endpoints(app: FastAPI) -> None:
router = APIRouter(
prefix="/traces/{trace_id}/feedback",
tags=["feedback"],
)
@router.put("/", status_code=status.HTTP_202_ACCEPTED)
def set_feedback(trace_id: str, input: EvaluationFeedbackRequest) -> Response:
trace = get_context().tracing_database.get(trace_id)
if trace is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
trace.feedback = input.feedback
get_context().tracing_database.update(trace_id, trace)
return Response(status_code=status.HTTP_202_ACCEPTED)
@router.get("/", status_code=status.HTTP_200_OK)
def get_feedback(trace_id: str) -> Any:
trace = get_context().tracing_database.get(trace_id)
if trace is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
return trace.feedback
@router.delete("/", status_code=status.HTTP_204_NO_CONTENT)
def delete_feedback(trace_id: str) -> Any:
trace = get_context().tracing_database.get(trace_id)
if trace is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
trace.feedback = None
get_context().tracing_database.update(trace_id, trace)
return Response(status_code=status.HTTP_204_NO_CONTENT)
app.include_router(router)

60
great_ai/src/great_ai/great_ai/deploy/routes/bootstrap_trace_endpoints.py Normal file
View file

@ -0,0 +1,60 @@
from pathlib import Path
from typing import List
from fastapi import APIRouter, FastAPI, HTTPException, Response, status
from fastapi.middleware.wsgi import WSGIMiddleware
from fastapi.responses import RedirectResponse
from fastapi.staticfiles import StaticFiles
from flask import Flask
from ...constants import METRICS_PATH
from ...context import get_context
from ...views import Query, Trace
PATH = Path(__file__).parent.resolve()
def bootstrap_trace_endpoints(app: FastAPI, dash_app: Flask) -> None:
app.mount(METRICS_PATH, WSGIMiddleware(dash_app))
@app.get("/", include_in_schema=False)
def redirect_to_entrypoint() -> RedirectResponse:
return RedirectResponse("/metrics")
app.mount(
"/assets",
StaticFiles(directory=PATH / "../../dashboard/assets"),
name="static",
)
router = APIRouter(
prefix="/traces",
tags=["traces"],
)
@router.post("/", status_code=status.HTTP_200_OK, response_model=List[Trace])
def query_traces(
query: Query,
skip: int = 0,
take: int = 100,
) -> List[Trace]:
return get_context().tracing_database.query(
conjunctive_filters=query.filter,
sort_by=query.sort,
skip=skip,
take=take,
)
@router.get("/{trace_id}", status_code=status.HTTP_200_OK, response_model=Trace)
def get_trace(trace_id: str) -> Trace:
result = get_context().tracing_database.get(trace_id)
if result is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
return result
@router.delete("/{trace_id}", status_code=status.HTTP_204_NO_CONTENT)
def delete_trace(trace_id: str) -> Response:
get_context().tracing_database.delete(trace_id)
return Response(status_code=status.HTTP_204_NO_CONTENT)
app.include_router(router)

1
great_ai/src/great_ai/great_ai/helper/freeze_arguments.py
View file

@ -22,4 +22,5 @@ def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]:
}
return func(*args, **kwargs)
wrapper.cache_info = func.cache_info # type: ignore
return wrapper

6
great_ai/src/great_ai/great_ai/models/load_model.py
View file

@ -2,17 +2,13 @@ from typing import Any, Optional, Tuple
from joblib import load
from great_ai.open_s3 import LargeFile
from ..context import get_context
def load_model(
key: str, version: Optional[int] = None, return_path: bool = False
) -> Tuple[Any, int]:
get_context() # will setup LargeFile if there was no config set
file = LargeFile(name=key, mode="rb", version=version)
file = get_context().large_file_implementation(name=key, mode="rb", version=version)
if return_path:
return file.get(), file.version

8
great_ai/src/great_ai/great_ai/models/save_model.py
View file

@ -3,17 +3,15 @@ from typing import Optional, Union
from joblib import dump
from great_ai.open_s3 import LargeFile
from ..context import get_context
def save_model(
model: Union[Path, str, object], key: str, keep_last_n: Optional[int] = None
) -> str:
get_context() # will setup LargeFile if there was no config set
file = LargeFile(name=key, mode="wb", keep_last_n=keep_last_n)
file = get_context().large_file_implementation(
name=key, mode="wb", keep_last_n=keep_last_n
)
if isinstance(model, Path) or isinstance(model, str):
file.push(model)

2
great_ai/src/great_ai/great_ai/models/use_model.py
View file

@ -2,7 +2,7 @@ from functools import wraps
from typing import Any, Callable, Dict, List, Literal, Union
from ..helper import assert_function_is_not_finalised, get_function_metadata_store
from ..tracing import TracingContext
from ..tracing.tracing_context import TracingContext
from ..views import Model
from .load_model import load_model

2
great_ai/src/great_ai/great_ai/parameters/log_metric.py
View file

@ -3,7 +3,7 @@ from typing import Any
from great_ai.great_ai.context.get_context import get_context
from ..tracing import TracingContext
from ..tracing.tracing_context import TracingContext
def log_metric(argument_name: str, value: Any) -> None:

2
great_ai/src/great_ai/great_ai/parameters/parameter.py
View file

@ -7,7 +7,7 @@ from ..helper import (
get_arguments,
get_function_metadata_store,
)
from ..tracing import TracingContext
from ..tracing.tracing_context import TracingContext
def parameter(

3
great_ai/src/great_ai/great_ai/persistence/__init__.py
View file

@ -1,3 +0,0 @@
from .mongodb_driver import MongoDbDriver
from .parallel_tinydb_driver import ParallelTinyDbDriver
from .persistence_driver import PersistenceDriver

5
great_ai/src/great_ai/great_ai/persistence/mongodb_driver.py
View file

@ -1,5 +0,0 @@
from .persistence_driver import PersistenceDriver
class MongoDbDriver(PersistenceDriver):
pass

38
great_ai/src/great_ai/great_ai/persistence/persistence_driver.py
View file

@ -1,38 +0,0 @@
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional
from ..views import Filter, SortBy, Trace
class PersistenceDriver(ABC):
is_threadsafe: bool
@abstractmethod
def save_trace(self, document: Trace) -> str:
pass
@abstractmethod
def add_feedback(self, id: str, evaluation: Any) -> None:
pass
@abstractmethod
def get_trace(self, id: str) -> Optional[Trace]:
pass
@abstractmethod
def get_traces(self) -> List[Trace]:
pass
@abstractmethod
def get_documents(self) -> List[Dict[str, Any]]:
pass
@abstractmethod
def query(
self,
conjunctive_filters: List[Filter],
sort_by: List[SortBy] = [],
skip: int = 0,
take: Optional[int] = None,
) -> List[Dict[str, Any]]:
pass

4
great_ai/src/great_ai/great_ai/tracing/__init__.py
View file

@ -1 +1,3 @@
from .tracing_context import TracingContext
from .mongodb_driver import MongoDbDriver
from .parallel_tinydb_driver import ParallelTinyDbDriver
from .tracing_database import TracingDatabase

5
great_ai/src/great_ai/great_ai/tracing/mongodb_driver.py Normal file
View file

@ -0,0 +1,5 @@
from .tracing_database import TracingDatabase
class MongoDbDriver(TracingDatabase):
pass

47
great_ai/src/great_ai/great_ai/persistence/parallel_tinydb_driver.py → great_ai/src/great_ai/great_ai/tracing/parallel_tinydb_driver.py
View file

@ -6,7 +6,7 @@ import pandas as pd
from tinydb import TinyDB
from ..views import Filter, SortBy, Trace
from .persistence_driver import PersistenceDriver
from .tracing_database import TracingDatabase
lock = Lock()
@ -14,47 +14,36 @@ lock = Lock()
operator_mapping = {"=": "eq", "!=": "ne", "<": "lt", "<=": "le", ">": "gt", ">=": "ge"}
class ParallelTinyDbDriver(PersistenceDriver):
class ParallelTinyDbDriver(TracingDatabase):
is_threadsafe = True
def __init__(self, path_to_db: Path) -> None:
super().__init__()
self._path_to_db = path_to_db
def save_trace(self, trace: Trace) -> str:
def save(self, trace: Trace) -> str:
return self._safe_execute(lambda db: db.insert(trace.dict()))
def add_feedback(self, id: str, evaluation: Any) -> None:
self._safe_execute(
lambda db: db.update(
fields={"evaluation": evaluation},
cond=lambda d: d["evaluation_id"] == id,
)
)
def get_trace(self, id: str) -> Optional[Trace]:
value = self._safe_execute(
lambda db: db.get(lambda d: d["evaluation_id"] == id)
)
def get(self, id: str) -> Optional[Trace]:
value = self._safe_execute(lambda db: db.get(lambda d: d["trace_id"] == id))
if value:
value = Trace.parse_obj(value)
return value
def get_traces(self) -> List[Trace]:
return self._safe_execute(lambda db: [Trace.parse_obj(t) for t in db.all()])
def get_documents(self) -> List[Dict[str, Any]]:
documents = self.get_traces()
return [d.to_flat_dict() for d in documents]
def query(
self,
conjunctive_filters: List[Filter],
sort_by: List[SortBy] = [],
skip: int = 0,
take: Optional[int] = None,
conjunctive_filters: List[Filter] = [],
sort_by: List[SortBy] = [],
) -> List[Dict[str, Any]]:
documents = self.get_documents()
documents = [
d.to_flat_dict()
for d in self._safe_execute(
lambda db: [Trace.parse_obj(t) for t in db.all()]
)
]
if not documents:
return []
@ -79,6 +68,14 @@ class ParallelTinyDbDriver(PersistenceDriver):
return result.to_dict("records")
def update(self, id: str, new_version: Trace) -> None:
self._safe_execute(
lambda db: db.update(new_version.dict(), lambda d: d["trace_id"] == id)
)
def delete(self, id: str) -> None:
self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] == id))
def _safe_execute(self, func: Callable[[TinyDB], Any]) -> Any:
with lock:
with TinyDB(self._path_to_db) as db:

4
great_ai/src/great_ai/great_ai/tracing/tracing_context.py
View file

@ -4,7 +4,7 @@ from datetime import datetime
from types import TracebackType
from typing import Any, DefaultDict, Dict, List, Literal, Optional, Type
from ..context import get_context
from ..context.get_context import get_context
from ..views import Model, Trace
@ -69,6 +69,6 @@ class TracingContext:
)
assert self._trace is not None
get_context().persistence.save_trace(self._trace)
get_context().tracing_database.save(self._trace)
return False

34
great_ai/src/great_ai/great_ai/tracing/tracing_database.py Normal file
View file

@ -0,0 +1,34 @@
from abc import ABC, abstractmethod
from typing import List, Optional
from ..views import Filter, SortBy, Trace
class TracingDatabase(ABC):
is_threadsafe: bool
@abstractmethod
def save(self, document: Trace) -> str:
pass
@abstractmethod
def get(self, id: str) -> Optional[Trace]:
pass
@abstractmethod
def query(
self,
skip: int = 0,
take: Optional[int] = None,
conjunctive_filters: List[Filter] = [],
sort_by: List[SortBy] = [],
) -> List[Trace]:
pass
@abstractmethod
def update(self, id: str, new_version: Trace) -> None:
pass
@abstractmethod
def delete(self, id: str) -> None:
pass

1
great_ai/src/great_ai/great_ai/views/__init__.py
View file

@ -1,4 +1,5 @@
from .api_metadata import ApiMetadata
from .cache_statistics import CacheStatistics
from .evaluation_feedback_request import EvaluationFeedbackRequest
from .filter import Filter
from .function_metadata import FunctionMetadata

3
great_ai/src/great_ai/great_ai/views/api_metadata.py
View file

@ -1,3 +1,5 @@
from typing import Any
from pydantic import BaseModel
@ -5,3 +7,4 @@ class ApiMetadata(BaseModel):
name: str
version: str
documentation: str
configuration: Any

8
great_ai/src/great_ai/great_ai/views/cache_statistics.py Normal file
View file

@ -0,0 +1,8 @@
from pydantic import BaseModel
class CacheStatistics(BaseModel):
hits: int
misses: int
size: int
max_size: int

2
great_ai/src/great_ai/great_ai/views/evaluation_feedback_request.py
View file

@ -4,4 +4,4 @@ from pydantic import BaseModel
class EvaluationFeedbackRequest(BaseModel):
evaluation: Any
feedback: Any

3
great_ai/src/great_ai/great_ai/views/health_check_response.py
View file

@ -1,5 +1,8 @@
from pydantic import BaseModel
from .cache_statistics import CacheStatistics
class HealthCheckResponse(BaseModel):
is_healthy: bool
cache_statistics: CacheStatistics

3
great_ai/src/great_ai/great_ai/views/query.py
View file

@ -9,8 +9,6 @@ from .sort_by import SortBy
class Query(BaseModel):
filter: List[Filter] = []
sort: List[SortBy] = []
skip: int = 0
take: int = 100
class Config:
schema_extra = {
@ -22,6 +20,5 @@ class Query(BaseModel):
{"column_id": "execution_time_ms", "direction": "asc"},
{"column_id": "id", "direction": "desc"},
],
"take": 10,
}
}

14
great_ai/src/great_ai/great_ai/views/trace.py
View file

@ -8,32 +8,30 @@ from .model import Model
class Trace(BaseModel):
evaluation_id: Optional[str]
trace_id: Optional[str]
created: str
execution_time_ms: float
logged_values: Dict[str, Any]
models: List[Model]
exception: Optional[str]
output: Any
evaluation: Any = None
feedback: Any = None
@validator("evaluation_id", always=True)
def validate_single_set(
cls, v: Optional[str], values: Dict[str, Any]
) -> Optional[str]:
@validator("trace_id", always=True)
def generate_id(cls, v: Optional[str], values: Dict[str, Any]) -> Optional[str]:
if not v:
return str(uuid4())
return v
def to_flat_dict(self) -> Dict[str, Any]:
return {
"id": self.evaluation_id,
"id": self.trace_id,
"created": self.created,
"execution_time_ms": self.execution_time_ms,
"models": ", ".join(f"{m.key}:{m.version}" for m in self.models),
"output": dumps(self.output),
"exception": self.exception or "null",
"evaluation": self.evaluation,
"feedback": self.feedback,
**self.logged_values,
}