mirror of
https://github.com/langchain-ai/kork.git
synced 2026-06-30 22:28:05 -04:00
Eugene Yurtsev
28 lines
780 B
Python
28 lines
780 B
Python
from kork.ast import ExternFunctionDef, ParamList
|
|
from kork.retrieval import SimpleContextRetriever
|
|
|
|
|
|
def foo() -> None:
|
|
"""Do nothing."""
|
|
|
|
|
|
def bar(x: int) -> int:
|
|
"""Add one to x."""
|
|
return x + 1
|
|
|
|
|
|
def test_simple_retriever() -> None:
|
|
"""Test simple retriever"""
|
|
external_func = ExternFunctionDef(
|
|
name="meow",
|
|
params=ParamList(params=[]),
|
|
return_type="int",
|
|
implementation=None,
|
|
)
|
|
simple_retriever = SimpleContextRetriever.from_functions([foo, bar, external_func])
|
|
external_func_defs = simple_retriever.retrieve("does not matter")
|
|
assert len(external_func_defs) == 3
|
|
assert external_func_defs[0].name == "foo"
|
|
assert external_func_defs[1].name == "bar"
|
|
assert external_func_defs[2].name == "meow"
|