Merge pull request #3 from biggraph/darabos-await-fmt

server/lynxscribe_ops.py

@@ -1,6 +1,7 @@
 """
 LynxScribe configuration and testing in LynxKite.
 """
+
 from lynxscribe.core.llm.base import get_llm_engine
 from lynxscribe.core.vector_store.base import get_vector_store
 from lynxscribe.common.config import load_config
@@ -9,7 +10,10 @@ from lynxscribe.components.rag.rag_graph import RAGGraph
 from lynxscribe.components.rag.knowledge_base_graph import PandasKnowledgeBaseGraph
 from lynxscribe.components.rag.rag_chatbot import Scenario, ScenarioSelector, RAGChatbot
 from lynxscribe.components.chat_processor.base import ChatProcessor
-from lynxscribe.components.chat_processor.processors import MaskTemplate, TruncateHistory
+from lynxscribe.components.chat_processor.processors import (
+    MaskTemplate,
+    TruncateHistory,
+)
 from lynxscribe.components.chat_api import ChatAPI, ChatAPIRequest, ChatAPIResponse
 
 from . import ops
@@ -17,187 +21,245 @@ import asyncio
 import json
 from .executors import one_by_one
 
-ENV = 'LynxScribe'
+ENV = "LynxScribe"
 one_by_one.register(ENV)
 op = ops.op_registration(ENV)
 output_on_top = ops.output_position(output="top")
 
+
 @output_on_top
 @op("Vector store")
-def vector_store(*, name='chromadb', collection_name='lynx'):
-  vector_store = get_vector_store(name=name, collection_name=collection_name)
-  return {'vector_store': vector_store}
+def vector_store(*, name="chromadb", collection_name="lynx"):
+    vector_store = get_vector_store(name=name, collection_name=collection_name)
+    return {"vector_store": vector_store}
+
 
 @output_on_top
 @op("LLM")
-def llm(*, name='openai'):
-  llm = get_llm_engine(name=name)
-  return {'llm': llm}
+def llm(*, name="openai"):
+    llm = get_llm_engine(name=name)
+    return {"llm": llm}
+
 
 @output_on_top
 @ops.input_position(llm="bottom")
 @op("Text embedder")
-def text_embedder(llm, *, model='text-embedding-ada-002'):
-  llm = llm[0]['llm']
-  text_embedder = TextEmbedder(llm=llm, model=model)
-  return {'text_embedder': text_embedder}
+def text_embedder(llm, *, model="text-embedding-ada-002"):
+    llm = llm[0]["llm"]
+    text_embedder = TextEmbedder(llm=llm, model=model)
+    return {"text_embedder": text_embedder}
+
 
 @output_on_top
 @ops.input_position(vector_store="bottom", text_embedder="bottom")
 @op("RAG graph")
 def rag_graph(vector_store, text_embedder):
-    vector_store = vector_store[0]['vector_store']
-    text_embedder = text_embedder[0]['text_embedder']
+    vector_store = vector_store[0]["vector_store"]
+    text_embedder = text_embedder[0]["text_embedder"]
     rag_graph = RAGGraph(
         PandasKnowledgeBaseGraph(vector_store=vector_store, text_embedder=text_embedder)
     )
-    return {'rag_graph': rag_graph}
+    return {"rag_graph": rag_graph}
+
 
 @output_on_top
 @op("Scenario selector")
-def scenario_selector(*, scenario_file: str, node_types='intent_cluster'):
-  scenarios = load_config(scenario_file)
-  node_types = [t.strip() for t in node_types.split(',')]
-  scenario_selector = ScenarioSelector(
-      scenarios=[Scenario(**scenario) for scenario in scenarios],
-      node_types=node_types,
-  )
-  return {'scenario_selector': scenario_selector}
+def scenario_selector(*, scenario_file: str, node_types="intent_cluster"):
+    scenarios = load_config(scenario_file)
+    node_types = [t.strip() for t in node_types.split(",")]
+    scenario_selector = ScenarioSelector(
+        scenarios=[Scenario(**scenario) for scenario in scenarios],
+        node_types=node_types,
+    )
+    return {"scenario_selector": scenario_selector}
+
 
 DEFAULT_NEGATIVE_ANSWER = "I'm sorry, but the data I've been trained on does not contain any information related to your question."
 
+
 @output_on_top
 @ops.input_position(rag_graph="bottom", scenario_selector="bottom", llm="bottom")
 @op("RAG chatbot")
 def rag_chatbot(
-    rag_graph, scenario_selector, llm, *,
+    rag_graph,
+    scenario_selector,
+    llm,
+    *,
     negative_answer=DEFAULT_NEGATIVE_ANSWER,
-    limits_by_type='{}',
-    strict_limits=True, max_results=5):
-  rag_graph = rag_graph[0]['rag_graph']
-  scenario_selector = scenario_selector[0]['scenario_selector']
-  llm = llm[0]['llm']
-  limits_by_type = json.loads(limits_by_type)
-  rag_chatbot = RAGChatbot(
-      rag_graph=rag_graph,
-      scenario_selector=scenario_selector,
-      llm=llm,
-      negative_answer=negative_answer,
-      limits_by_type=limits_by_type,
-      strict_limits=strict_limits,
-      max_results=max_results,
-  )
-  return {'chatbot': rag_chatbot}
+    limits_by_type="{}",
+    strict_limits=True,
+    max_results=5,
+):
+    rag_graph = rag_graph[0]["rag_graph"]
+    scenario_selector = scenario_selector[0]["scenario_selector"]
+    llm = llm[0]["llm"]
+    limits_by_type = json.loads(limits_by_type)
+    rag_chatbot = RAGChatbot(
+        rag_graph=rag_graph,
+        scenario_selector=scenario_selector,
+        llm=llm,
+        negative_answer=negative_answer,
+        limits_by_type=limits_by_type,
+        strict_limits=strict_limits,
+        max_results=max_results,
+    )
+    return {"chatbot": rag_chatbot}
+
 
 @output_on_top
 @ops.input_position(processor="bottom")
 @op("Chat processor")
 def chat_processor(processor, *, _ctx: one_by_one.Context):
-  cfg = _ctx.last_result or {'question_processors': [], 'answer_processors': [], 'masks': []}
-  for f in ['question_processor', 'answer_processor', 'mask']:
-    if f in processor:
-      cfg[f + 's'].append(processor[f])
-  question_processors = cfg['question_processors'][:]
-  answer_processors = cfg['answer_processors'][:]
-  masking_templates = {}
-  for mask in cfg['masks']:
-    masking_templates[mask['name']] = mask
-  if masking_templates:
-    question_processors.append(MaskTemplate(masking_templates=masking_templates))
-    answer_processors.append(MaskTemplate(masking_templates=masking_templates))
-  chat_processor = ChatProcessor(question_processors=question_processors, answer_processors=answer_processors)
-  return {'chat_processor': chat_processor, **cfg}
+    cfg = _ctx.last_result or {
+        "question_processors": [],
+        "answer_processors": [],
+        "masks": [],
+    }
+    for f in ["question_processor", "answer_processor", "mask"]:
+        if f in processor:
+            cfg[f + "s"].append(processor[f])
+    question_processors = cfg["question_processors"][:]
+    answer_processors = cfg["answer_processors"][:]
+    masking_templates = {}
+    for mask in cfg["masks"]:
+        masking_templates[mask["name"]] = mask
+    if masking_templates:
+        question_processors.append(MaskTemplate(masking_templates=masking_templates))
+        answer_processors.append(MaskTemplate(masking_templates=masking_templates))
+    chat_processor = ChatProcessor(
+        question_processors=question_processors, answer_processors=answer_processors
+    )
+    return {"chat_processor": chat_processor, **cfg}
+
 
 @output_on_top
 @op("Truncate history")
-def truncate_history(*, max_tokens=10000, language='English'):
-  return {'question_processor': TruncateHistory(max_tokens=max_tokens, language=language.lower())}
+def truncate_history(*, max_tokens=10000, language="English"):
+    return {
+        "question_processor": TruncateHistory(
+            max_tokens=max_tokens, language=language.lower()
+        )
+    }
+
 
 @output_on_top
 @op("Mask")
-def mask(*, name='', regex='', exceptions='', mask_pattern=''):
-  exceptions = [e.strip() for e in exceptions.split(',') if e.strip()]
-  return {'mask': {'name': name, 'regex': regex, 'exceptions': exceptions, 'mask_pattern': mask_pattern}}
+def mask(*, name="", regex="", exceptions="", mask_pattern=""):
+    exceptions = [e.strip() for e in exceptions.split(",") if e.strip()]
+    return {
+        "mask": {
+            "name": name,
+            "regex": regex,
+            "exceptions": exceptions,
+            "mask_pattern": mask_pattern,
+        }
+    }
+
 
 @ops.input_position(chat_api="bottom")
 @op("Test Chat API")
-def test_chat_api(message, chat_api, *, show_details=False):
-  chat_api = chat_api[0]['chat_api']
-  request = ChatAPIRequest(session_id="b43215a0-428f-11ef-9454-0242ac120002", question=message['text'], history=[])
-  response = asyncio.run(chat_api.answer(request))
-  if show_details:
-    return {**response.__dict__}
-  else:
-    return {'answer': response.answer}
+async def test_chat_api(message, chat_api, *, show_details=False):
+    chat_api = chat_api[0]["chat_api"]
+    request = ChatAPIRequest(
+        session_id="b43215a0-428f-11ef-9454-0242ac120002",
+        question=message["text"],
+        history=[],
+    )
+    response = await chat_api.answer(request)
+    if show_details:
+        return {**response.__dict__}
+    else:
+        return {"answer": response.answer}
+
 
 @op("Input chat")
 def input_chat(*, chat: str):
-  return {'text': chat}
+    return {"text": chat}
+
 
 @output_on_top
 @ops.input_position(chatbot="bottom", chat_processor="bottom", knowledge_base="bottom")
 @op("Chat API")
-def chat_api(chatbot, chat_processor, knowledge_base, *, model='gpt-4o-mini'):
-  chatbot = chatbot[0]['chatbot']
-  chat_processor = chat_processor[0]['chat_processor']
-  knowledge_base = knowledge_base[0]
-  c = ChatAPI(
-      chatbot=chatbot,
-      chat_processor=chat_processor,
-      model=model,
-  )
-  if knowledge_base:
-    c.chatbot.rag_graph.kg_base.load_v1_knowledge_base(**knowledge_base)
-    c.chatbot.scenario_selector.check_compatibility(c.chatbot.rag_graph)
-  return {'chat_api': c}
+def chat_api(chatbot, chat_processor, knowledge_base, *, model="gpt-4o-mini"):
+    chatbot = chatbot[0]["chatbot"]
+    chat_processor = chat_processor[0]["chat_processor"]
+    knowledge_base = knowledge_base[0]
+    c = ChatAPI(
+        chatbot=chatbot,
+        chat_processor=chat_processor,
+        model=model,
+    )
+    if knowledge_base:
+        c.chatbot.rag_graph.kg_base.load_v1_knowledge_base(**knowledge_base)
+        c.chatbot.scenario_selector.check_compatibility(c.chatbot.rag_graph)
+    return {"chat_api": c}
+
 
 @output_on_top
 @op("Knowledge base")
-def knowledge_base(*, nodes_path='nodes.pickle', edges_path='edges.pickle', template_cluster_path='tempclusters.pickle'):
-    return {'nodes_path': nodes_path, 'edges_path': edges_path, 'template_cluster_path': template_cluster_path}
+def knowledge_base(
+    *,
+    nodes_path="nodes.pickle",
+    edges_path="edges.pickle",
+    template_cluster_path="tempclusters.pickle",
+):
+    return {
+        "nodes_path": nodes_path,
+        "edges_path": edges_path,
+        "template_cluster_path": template_cluster_path,
+    }
+
 
 @op("View", view="table_view")
 def view(input):
-  columns = [str(c) for c in input.keys() if not str(c).startswith('_')]
-  v = {
-    'dataframes': { 'df': {
-      'columns': columns,
-      'data': [[input[c] for c in columns]],
-    }}
-  }
-  return v
+    columns = [str(c) for c in input.keys() if not str(c).startswith("_")]
+    v = {
+        "dataframes": {
+            "df": {
+                "columns": columns,
+                "data": [[input[c] for c in columns]],
+            }
+        }
+    }
+    return v
+
 
 async def api_service(request):
-  '''
-  Serves a chat endpoint that matches LynxScribe's interface.
-  To access it you need to add the "module" and "workspace"
-  parameters.
-  The workspace must contain exactly one "Chat API" node.
-
-    curl -X POST ${LYNXKITE_URL}/api/service \
-      -H "Content-Type: application/json" \
-      -d '{
-        "module": "server.lynxscribe_ops",
-        "workspace": "LynxScribe demo",
-        "session_id": "b43215a0-428f-11ef-9454-0242ac120002",
-        "question": "what does the fox say",
-        "history": [],
-        "user_id": "x",
-        "meta_inputs": {}
-      }'
-  '''
-  import pathlib
-  from . import workspace
-  DATA_PATH = pathlib.Path.cwd() / 'data'
-  path = DATA_PATH / request['workspace']
-  assert path.is_relative_to(DATA_PATH)
-  assert path.exists(), f'Workspace {path} does not exist'
-  ws = workspace.load(path)
-  contexts = ops.EXECUTORS[ENV](ws)
-  nodes = [op for op in ws.nodes if op.data.title == 'Chat API']
-  [node] = nodes
-  context = contexts[node.id]
-  chat_api = context.last_result['chat_api']
-  request = ChatAPIRequest(session_id=request['session_id'], question=request['question'], history=request['history'])
-  response = await chat_api.answer(request)
-  return response
+    """
+    Serves a chat endpoint that matches LynxScribe's interface.
+    To access it you need to add the "module" and "workspace"
+    parameters.
+    The workspace must contain exactly one "Chat API" node.
+
+      curl -X POST ${LYNXKITE_URL}/api/service \
+        -H "Content-Type: application/json" \
+        -d '{
+          "module": "server.lynxscribe_ops",
+          "workspace": "LynxScribe demo",
+          "session_id": "b43215a0-428f-11ef-9454-0242ac120002",
+          "question": "what does the fox say",
+          "history": [],
+          "user_id": "x",
+          "meta_inputs": {}
+        }'
+  """
+    import pathlib
+    from . import workspace
+
+    DATA_PATH = pathlib.Path.cwd() / "data"
+    path = DATA_PATH / request["workspace"]
+    assert path.is_relative_to(DATA_PATH)
+    assert path.exists(), f"Workspace {path} does not exist"
+    ws = workspace.load(path)
+    contexts = ops.EXECUTORS[ENV](ws)
+    nodes = [op for op in ws.nodes if op.data.title == "Chat API"]
+    [node] = nodes
+    context = contexts[node.id]
+    chat_api = context.last_result["chat_api"]
+    request = ChatAPIRequest(
+        session_id=request["session_id"],
+        question=request["question"],
+        history=request["history"],
+    )
+    response = await chat_api.answer(request)
+    return response

server/main.py

@@ -9,30 +9,34 @@ from . import workspace
 here = pathlib.Path(__file__).parent
 lynxkite_modules = {}
 for _, name, _ in pkgutil.iter_modules([str(here)]):
-    if name.endswith('_ops') and not name.startswith('test_'):
-        print(f'Importing {name}')
-        name = f'server.{name}'
+    if name.endswith("_ops") and not name.startswith("test_"):
+        print(f"Importing {name}")
+        name = f"server.{name}"
         lynxkite_modules[name] = importlib.import_module(name)
 
 app = fastapi.FastAPI()
 
 
+
 @app.get("/api/catalog")
 def get_catalog():
     return {
         k: {op.name: op.model_dump() for op in v.values()}
-        for k, v in ops.CATALOGS.items()}
+        for k, v in ops.CATALOGS.items()
+    }
 
 
 class SaveRequest(workspace.BaseConfig):
     path: str
     ws: workspace.Workspace
 
+
 def save(req: SaveRequest):
     path = DATA_PATH / req.path
     assert path.is_relative_to(DATA_PATH)
     workspace.save(req.ws, path)
 
+
 @app.post("/api/save")
 def save_and_execute(req: SaveRequest):
     save(req)
@@ -40,6 +44,7 @@ def save_and_execute(req: SaveRequest):
     save(req)
     return req.ws
 
+
 @app.get("/api/load")
 def load(path: str):
     path = DATA_PATH / path
@@ -48,31 +53,41 @@ def load(path: str):
         return workspace.Workspace()
     return workspace.load(path)
 
-DATA_PATH = pathlib.Path.cwd() / 'data'
+
+DATA_PATH = pathlib.Path.cwd() / "data"
+
 
 @dataclasses.dataclass(order=True)
 class DirectoryEntry:
     name: str
     type: str
 
+
 @app.get("/api/dir/list")
 def list_dir(path: str):
     path = DATA_PATH / path
     assert path.is_relative_to(DATA_PATH)
-    return sorted([
-        DirectoryEntry(p.relative_to(DATA_PATH), 'directory' if p.is_dir() else 'workspace')
-        for p in path.iterdir()])
+    return sorted(
+        [
+            DirectoryEntry(
+                p.relative_to(DATA_PATH), "directory" if p.is_dir() else "workspace"
+            )
+            for p in path.iterdir()
+        ]
+    )
+
 
 @app.post("/api/dir/mkdir")
 def make_dir(req: dict):
-    path = DATA_PATH / req['path']
+    path = DATA_PATH / req["path"]
     assert path.is_relative_to(DATA_PATH)
     assert not path.exists()
     path.mkdir()
     return list_dir(path.parent)
 
+
 @app.post("/api/service")
 async def service(req: dict):
-    '''Executors can provide extra HTTP APIs through the /api/service endpoint.'''
-    module = lynxkite_modules[req['module']]
+    """Executors can provide extra HTTP APIs through the /api/service endpoint."""
+    module = lynxkite_modules[req["module"]]
     return await module.api_service(req)