Compare commits
4 Commits
370301cd15
...
2c8e93dcde
| Author | SHA1 | Date |
|---|---|---|
|
|
2c8e93dcde | |
|
|
a5f33c88d1 | |
|
|
00038b679a | |
|
|
2c63a030e5 |
|
|
@ -66,4 +66,7 @@ docker-compose.override.yml
|
|||
node_modules/
|
||||
npm-debug.log*
|
||||
yarn-debug.log*
|
||||
yarn-error.log*
|
||||
yarn-error.log*
|
||||
|
||||
# db
|
||||
chroma_db
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
# 主模型
|
||||
MAIN_HTTP_URL=http://localhost:8022/v1/chat/completions
|
||||
MAIN_HTTP_KEY=sk-local-827ccb0eea8a706c4c34a16891f84e7b
|
||||
MAIN_MODEL_NAME=Qwen2.5
|
||||
# Embedding
|
||||
EMBEDDING_HTTP_URL=http://localhost:8023/v1/embeddings
|
||||
EMBEDDING_HTTP_KEY=sk-local-827ccb0eea8a706c4c34a16891f84e7b
|
||||
EMBEDDING_MODEL_NAME=Qwen3-Embedding
|
||||
32
README.md
32
README.md
|
|
@ -0,0 +1,32 @@
|
|||
## 目录结构
|
||||
|
||||
```
|
||||
sample-ai/
|
||||
│
|
||||
├── requirements.txt
|
||||
├── README.md
|
||||
├── setup.py # 可选,打包/安装用
|
||||
│
|
||||
├── src/
|
||||
│ ├── sample_ai/
|
||||
│ │ ├── __init__.py
|
||||
│ │ ├── main.py # 项目主入口(如有 CLI 或服务)
|
||||
│ │ ├── nodes.py # 流程节点相关
|
||||
│ │ ├── vector_store.py # 向量数据库/知识库管理
|
||||
│ │ ├── utils.py # 工具函数(如 embedding、分块等)
|
||||
│ │ ├── config.py # 配置管理(如环境变量、常量等)
|
||||
│ │ └── data/
|
||||
│ │ ├── __init__.py
|
||||
│ │ └── ... # 数据处理相关(如 chunk、预处理等)
|
||||
│ │
|
||||
│ └── tests/
|
||||
│ ├── __init__.py
|
||||
│ ├── test_nodes.py
|
||||
│ ├── test_vector_store.py
|
||||
│ └── test_utils.py
|
||||
│
|
||||
├── scripts/
|
||||
│ └── run_server.sh # 启动脚本等
|
||||
│
|
||||
└── chroma_db/ # Chroma 本地数据库(自动生成,可忽略版本控制)
|
||||
```
|
||||
|
|
@ -1 +1,5 @@
|
|||
httpx==0.28.1
|
||||
httpx==0.28.1
|
||||
chromadb==1.3.4
|
||||
python-dotenv==1.2.1
|
||||
pytest==9.0.1
|
||||
pytest-asyncio==1.3.0
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
from sample_ai.pocketflow import AsyncFlow
|
||||
from sample_ai.nodes import ChunkDocumentsNode
|
||||
|
||||
def create_online_flow():
|
||||
chunk_node = ChunkDocumentsNode()
|
||||
flow = AsyncFlow(chunk_node)
|
||||
return flow
|
||||
|
|
@ -0,0 +1,108 @@
|
|||
from sample_ai.pocketflow import AsyncBatchNode, AsyncFlow, Node
|
||||
from sample_ai.utils import fixed_size_chunk
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import chromadb
|
||||
|
||||
# Nodes for the offline flow
|
||||
class ChunkDocumentsNode(AsyncBatchNode):
|
||||
async def prep_async(self, shared):
|
||||
"""Read texts from shared store"""
|
||||
return shared["texts"]
|
||||
|
||||
async def exec_async(self, text):
|
||||
"""Chunk a single text into smaller pieces"""
|
||||
# 先将所有制表符等替换为一个空格
|
||||
text = re.sub(r'[ \t]+', ' ', text)
|
||||
# 再将多个空格替换为一个空格
|
||||
text = re.sub(r' +', ' ', text)
|
||||
# 去除首尾空格
|
||||
text = text.strip()
|
||||
return fixed_size_chunk(text, chunk_size=500, overlap=50)
|
||||
|
||||
async def post_async(self, shared, prep_res, exec_res_list):
|
||||
"""Store chunked texts in the shared store"""
|
||||
# Flatten the list of lists into a single list of chunks
|
||||
all_chunks = []
|
||||
for chunks in exec_res_list:
|
||||
all_chunks.extend(chunks)
|
||||
|
||||
# Replace the original texts with the flat list of chunks
|
||||
print(f"all_chunks: {json.dumps(all_chunks, indent=2)}")
|
||||
shared["texts"] = all_chunks
|
||||
|
||||
print(f"✅ Created {len(all_chunks)} chunks from {len(prep_res)} documents")
|
||||
return "default"
|
||||
|
||||
|
||||
class EmbedDocumentsNode(AsyncBatchNode):
|
||||
async def prep_async(self, shared):
|
||||
return shared["texts"]
|
||||
|
||||
async def exec_async(self, text):
|
||||
return fixed_size_chunk(text, chunk_size=500, overlap=50)
|
||||
|
||||
async def post_async(self, shared, prep_res, exec_res_list):
|
||||
return "default"
|
||||
|
||||
|
||||
class CreateIndexNode(Node):
|
||||
def prep(self, shared):
|
||||
"""Get embeddings from shared store"""
|
||||
return shared["embeddings"]
|
||||
|
||||
def exec(self, embeddings):
|
||||
"""Create FAISS index and add embeddings"""
|
||||
return []
|
||||
|
||||
def post(self, shared, prep_res, exec_res):
|
||||
"""Store the index in shared store"""
|
||||
shared["index"] = exec_res
|
||||
print(f"✅ Index created with {exec_res.ntotal} vectors")
|
||||
return "default"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
async def _do():
|
||||
shared = {
|
||||
"texts": [
|
||||
# PocketFlow framework
|
||||
"""Pocket Flow is a 100-line minimalist LLM framework
|
||||
Lightweight: Just 100 lines. Zero bloat, zero dependencies, zero vendor lock-in.
|
||||
Expressive: Everything you love—(Multi-)Agents, Workflow, RAG, and more.
|
||||
Agentic Coding: Let AI Agents (e.g., Cursor AI) build Agents—10x productivity boost!
|
||||
To install, pip install pocketflow or just copy the source code (only 100 lines).""",
|
||||
# Fictional medical device
|
||||
"""NeurAlign M7 is a revolutionary non-invasive neural alignment device.
|
||||
Targeted magnetic resonance technology increases neuroplasticity in specific brain regions.
|
||||
Clinical trials showed 72% improvement in PTSD treatment outcomes.
|
||||
Developed by Cortex Medical in 2024 as an adjunct to standard cognitive therapy.
|
||||
Portable design allows for in-home use with remote practitioner monitoring.""",
|
||||
# Made-up historical event
|
||||
"""The Velvet Revolution of Caldonia (1967-1968) ended Generalissimo Verak's 40-year rule.
|
||||
Led by poet Eliza Markovian through underground literary societies.
|
||||
Culminated in the Great Silence Protest with 300,000 silent protesters.
|
||||
First democratic elections held in March 1968 with 94% voter turnout.
|
||||
Became a model for non-violent political transitions in neighboring regions.""",
|
||||
# Fictional technology
|
||||
"""Q-Mesh is QuantumLeap Technologies' instantaneous data synchronization protocol.
|
||||
Utilizes directed acyclic graph consensus for 500,000 transactions per second.
|
||||
Consumes 95% less energy than traditional blockchain systems.
|
||||
Adopted by three central banks for secure financial data transfer.
|
||||
Released in February 2024 after five years of development in stealth mode.""",
|
||||
# Made-up scientific research
|
||||
"""Harlow Institute's Mycelium Strain HI-271 removes 99.7% of PFAS from contaminated soil.
|
||||
Engineered fungi create symbiotic relationships with native soil bacteria.
|
||||
Breaks down "forever chemicals" into non-toxic compounds within 60 days.
|
||||
Field tests successfully remediated previously permanently contaminated industrial sites.
|
||||
Deployment costs 80% less than traditional chemical extraction methods.""",
|
||||
]
|
||||
}
|
||||
|
||||
chunk_node = ChunkDocumentsNode()
|
||||
flow = AsyncFlow(chunk_node)
|
||||
await flow.run_async(shared)
|
||||
|
||||
asyncio.run(_do())
|
||||
|
|
@ -0,0 +1,197 @@
|
|||
import asyncio, warnings, copy, time
|
||||
|
||||
|
||||
class BaseNode:
|
||||
def __init__(self):
|
||||
self.params, self.successors = {}, {}
|
||||
|
||||
def set_params(self, params):
|
||||
self.params = params
|
||||
|
||||
def next(self, node, action="default"):
|
||||
if action in self.successors:
|
||||
warnings.warn(f"Overwriting successor for action '{action}'")
|
||||
self.successors[action] = node
|
||||
return node
|
||||
|
||||
def prep(self, shared):
|
||||
pass
|
||||
|
||||
def exec(self, prep_res):
|
||||
pass
|
||||
|
||||
def post(self, shared, prep_res, exec_res):
|
||||
pass
|
||||
|
||||
def _exec(self, prep_res):
|
||||
return self.exec(prep_res)
|
||||
|
||||
def _run(self, shared):
|
||||
p = self.prep(shared)
|
||||
e = self._exec(p)
|
||||
return self.post(shared, p, e)
|
||||
|
||||
def run(self, shared):
|
||||
if self.successors:
|
||||
warnings.warn("Node won't run successors. Use Flow.")
|
||||
return self._run(shared)
|
||||
|
||||
def __rshift__(self, other):
|
||||
return self.next(other)
|
||||
|
||||
def __sub__(self, action):
|
||||
if isinstance(action, str):
|
||||
return _ConditionalTransition(self, action)
|
||||
raise TypeError("Action must be a string")
|
||||
|
||||
|
||||
class _ConditionalTransition:
|
||||
def __init__(self, src, action):
|
||||
self.src, self.action = src, action
|
||||
|
||||
def __rshift__(self, tgt):
|
||||
return self.src.next(tgt, self.action)
|
||||
|
||||
|
||||
class Node(BaseNode):
|
||||
def __init__(self, max_retries=1, wait=0):
|
||||
super().__init__()
|
||||
self.max_retries, self.wait = max_retries, wait
|
||||
|
||||
def exec_fallback(self, prep_res, exc):
|
||||
raise exc
|
||||
|
||||
def _exec(self, prep_res):
|
||||
for self.cur_retry in range(self.max_retries):
|
||||
try:
|
||||
return self.exec(prep_res)
|
||||
except Exception as e:
|
||||
if self.cur_retry == self.max_retries - 1:
|
||||
return self.exec_fallback(prep_res, e)
|
||||
if self.wait > 0:
|
||||
time.sleep(self.wait)
|
||||
|
||||
|
||||
class BatchNode(Node):
|
||||
def _exec(self, items):
|
||||
return [super(BatchNode, self)._exec(i) for i in (items or [])]
|
||||
|
||||
|
||||
class Flow(BaseNode):
|
||||
def __init__(self, start=None):
|
||||
super().__init__()
|
||||
self.start_node = start
|
||||
|
||||
def start(self, start):
|
||||
self.start_node = start
|
||||
return start
|
||||
|
||||
def get_next_node(self, curr, action):
|
||||
nxt = curr.successors.get(action or "default")
|
||||
if not nxt and curr.successors:
|
||||
warnings.warn(f"Flow ends: '{action}' not found in {list(curr.successors)}")
|
||||
return nxt
|
||||
|
||||
def _orch(self, shared, params=None):
|
||||
curr, p, last_action = copy.copy(self.start_node), (params or {**self.params}), None
|
||||
while curr:
|
||||
curr.set_params(p)
|
||||
last_action = curr._run(shared)
|
||||
curr = copy.copy(self.get_next_node(curr, last_action))
|
||||
return last_action
|
||||
|
||||
def _run(self, shared):
|
||||
p = self.prep(shared)
|
||||
o = self._orch(shared)
|
||||
return self.post(shared, p, o)
|
||||
|
||||
def post(self, shared, prep_res, exec_res):
|
||||
return exec_res
|
||||
|
||||
|
||||
class BatchFlow(Flow):
|
||||
def _run(self, shared):
|
||||
pr = self.prep(shared) or []
|
||||
for bp in pr:
|
||||
self._orch(shared, {**self.params, **bp})
|
||||
return self.post(shared, pr, None)
|
||||
|
||||
|
||||
class AsyncNode(Node):
|
||||
async def prep_async(self, shared):
|
||||
pass
|
||||
|
||||
async def exec_async(self, prep_res):
|
||||
pass
|
||||
|
||||
async def exec_fallback_async(self, prep_res, exc):
|
||||
raise exc
|
||||
|
||||
async def post_async(self, shared, prep_res, exec_res):
|
||||
pass
|
||||
|
||||
async def _exec(self, prep_res):
|
||||
for self.cur_retry in range(self.max_retries):
|
||||
try:
|
||||
return await self.exec_async(prep_res)
|
||||
except Exception as e:
|
||||
if self.cur_retry == self.max_retries - 1:
|
||||
return await self.exec_fallback_async(prep_res, e)
|
||||
if self.wait > 0:
|
||||
await asyncio.sleep(self.wait)
|
||||
|
||||
async def run_async(self, shared):
|
||||
if self.successors:
|
||||
warnings.warn("Node won't run successors. Use AsyncFlow.")
|
||||
return await self._run_async(shared)
|
||||
|
||||
async def _run_async(self, shared):
|
||||
p = await self.prep_async(shared)
|
||||
e = await self._exec(p)
|
||||
return await self.post_async(shared, p, e)
|
||||
|
||||
def _run(self, shared):
|
||||
raise RuntimeError("Use run_async.")
|
||||
|
||||
|
||||
class AsyncBatchNode(AsyncNode, BatchNode):
|
||||
async def _exec(self, items):
|
||||
return [await super(AsyncBatchNode, self)._exec(i) for i in items]
|
||||
|
||||
|
||||
class AsyncParallelBatchNode(AsyncNode, BatchNode):
|
||||
async def _exec(self, items):
|
||||
return await asyncio.gather(*(super(AsyncParallelBatchNode, self)._exec(i) for i in items))
|
||||
|
||||
|
||||
class AsyncFlow(Flow, AsyncNode):
|
||||
async def _orch_async(self, shared, params=None):
|
||||
curr, p, last_action = copy.copy(self.start_node), (params or {**self.params}), None
|
||||
while curr:
|
||||
curr.set_params(p)
|
||||
last_action = await curr._run_async(shared) if isinstance(curr, AsyncNode) else curr._run(shared)
|
||||
curr = copy.copy(self.get_next_node(curr, last_action))
|
||||
return last_action
|
||||
|
||||
async def _run_async(self, shared):
|
||||
p = await self.prep_async(shared)
|
||||
o = await self._orch_async(shared)
|
||||
return await self.post_async(shared, p, o)
|
||||
|
||||
async def post_async(self, shared, prep_res, exec_res):
|
||||
return exec_res
|
||||
|
||||
|
||||
class AsyncBatchFlow(AsyncFlow, BatchFlow):
|
||||
async def _run_async(self, shared):
|
||||
pr = await self.prep_async(shared) or []
|
||||
for bp in pr:
|
||||
await self._orch_async(shared, {**self.params, **bp})
|
||||
return await self.post_async(shared, pr, None)
|
||||
|
||||
|
||||
class AsyncParallelBatchFlow(AsyncFlow, BatchFlow):
|
||||
async def _run_async(self, shared):
|
||||
pr = await self.prep_async(shared) or []
|
||||
await asyncio.gather(*(self._orch_async(shared, {**self.params, **bp}) for bp in pr))
|
||||
return await self.post_async(shared, pr, None)
|
||||
|
|
@ -1,18 +1,18 @@
|
|||
from venv import logger
|
||||
import httpx
|
||||
import asyncio
|
||||
import os
|
||||
import json
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
|
||||
VLLM_HTTP_KEY = os.environ.get("VLLM_HTTP_KEY")
|
||||
VLLM_HTTP_URL = os.environ.get("VLLM_HTTP_URL")
|
||||
VLLM_CHAT_MODEL_NAME = os.environ.get("VLLM_CHAT_MODEL_NAME")
|
||||
MAIN_HTTP_URL = os.environ.get("MAIN_HTTP_URL")
|
||||
MAIN_HTTP_KEY = os.environ.get("MAIN_HTTP_KEY")
|
||||
MAIN_MODEL_NAME = os.environ.get("MAIN_MODEL_NAME")
|
||||
|
||||
|
||||
def _headers():
|
||||
h = {"Content-Type": "application/json"}
|
||||
if VLLM_HTTP_KEY:
|
||||
h["Authorization"] = f"Bearer {VLLM_HTTP_KEY}"
|
||||
return h
|
||||
EMBEDDING_HTTP_URL = os.environ.get("EMBEDDING_HTTP_URL")
|
||||
EMBEDDING_HTTP_KEY = os.environ.get("EMBEDDING_HTTP_KEY")
|
||||
EMBEDDING_MODEL_NAME = os.environ.get("EMBEDDING_MODEL_NAME")
|
||||
|
||||
|
||||
def _merge_messages_and_prompt(messages: list[dict] = [], prompt: str = ""):
|
||||
|
|
@ -45,18 +45,32 @@ def _get_content(resp: dict) -> str:
|
|||
return ""
|
||||
|
||||
|
||||
async def call_llm(prompt: str, messages = [], max_tokens: int = 512, temperature: float = 0.0, timeout: int = 30):
|
||||
def fixed_size_chunk(text, chunk_size=2000, overlap=50):
|
||||
chunks = []
|
||||
start = 0
|
||||
while start < len(text):
|
||||
end = start + chunk_size
|
||||
chunks.append(text[max(0, start - overlap) : min(len(text), end + overlap)])
|
||||
start += chunk_size
|
||||
return chunks
|
||||
|
||||
|
||||
async def call_llm(prompt: str, messages=[], max_tokens: int = 512, temperature: float = 0.0, timeout: int = 30):
|
||||
"""简单对话"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
url = VLLM_HTTP_URL.rstrip("/") + "/v1/chat/completions"
|
||||
url = MAIN_HTTP_URL
|
||||
body = {
|
||||
"model": VLLM_CHAT_MODEL_NAME,
|
||||
"model": MAIN_MODEL_NAME,
|
||||
"messages": _merge_messages_and_prompt(messages, prompt),
|
||||
"max_tokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
}
|
||||
res = await client.post(url, headers=_headers(), json=body)
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {MAIN_HTTP_KEY}",
|
||||
}
|
||||
res = await client.post(url, headers=headers, json=body)
|
||||
res.raise_for_status()
|
||||
data = res.json()
|
||||
return _get_content(data)
|
||||
|
|
@ -65,29 +79,33 @@ async def call_llm(prompt: str, messages = [], max_tokens: int = 512, temperatur
|
|||
return ""
|
||||
|
||||
|
||||
async def call_llm_stream(prompt: str, messages = [], max_tokens: int = 512, temperature: float = 0.0, timeout: int = 60):
|
||||
async def call_llm_stream(prompt: str, messages=[], max_tokens: int = 512, temperature: float = 0.0, timeout: int = 60):
|
||||
"""
|
||||
流式对话
|
||||
- 使用: async for chunk call_llm_stream("prompt"):
|
||||
"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
url = (VLLM_HTTP_URL or "").rstrip("/") + "/v1/chat/completions"
|
||||
url = MAIN_HTTP_URL
|
||||
body = {
|
||||
"model": VLLM_CHAT_MODEL_NAME,
|
||||
"model": MAIN_MODEL_NAME,
|
||||
"messages": _merge_messages_and_prompt(messages or [], prompt),
|
||||
"max_tokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
"stream": True,
|
||||
}
|
||||
async with client.stream("POST", url, headers=_headers(), json=body) as resp:
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {MAIN_HTTP_KEY}",
|
||||
}
|
||||
async with client.stream("POST", url, headers=headers, json=body) as resp:
|
||||
resp.raise_for_status()
|
||||
async for raw_line in resp.aiter_lines():
|
||||
if not raw_line:
|
||||
continue
|
||||
line = raw_line.strip()
|
||||
if line.startswith("data:"):
|
||||
payload = line[len("data:"):].strip()
|
||||
payload = line[len("data:") :].strip()
|
||||
data = {}
|
||||
if payload in ("[DONE]", ""):
|
||||
break
|
||||
|
|
@ -105,31 +123,62 @@ async def call_llm_stream(prompt: str, messages = [], max_tokens: int = 512, tem
|
|||
return
|
||||
|
||||
|
||||
async def get_embedding(text, timeout: int = 30):
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
url = EMBEDDING_HTTP_URL
|
||||
body = {
|
||||
"model": EMBEDDING_MODEL_NAME,
|
||||
"input": [text],
|
||||
}
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {EMBEDDING_HTTP_KEY}",
|
||||
}
|
||||
res = await client.post(url, headers=headers, json=body)
|
||||
res.raise_for_status()
|
||||
data = res.json()
|
||||
return data["data"][0]["embedding"]
|
||||
except Exception as e:
|
||||
print(f"call_llm[ERROR]: {e}")
|
||||
return []
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 设置
|
||||
os.environ.setdefault("VLLM_HTTP_URL", "http://localhost:8022")
|
||||
os.environ.setdefault("VLLM_HTTP_KEY", "sk-local-827ccb0eea8a706c4c34a16891f84e7b")
|
||||
os.environ.setdefault("VLLM_CHAT_MODEL_NAME", "Qwen2.5-0.5B-Instruct")
|
||||
os.environ.setdefault("MAIN_HTTP_URL", "http://localhost:8022/v1/chat/completions")
|
||||
os.environ.setdefault("MAIN_HTTP_KEY", "sk-local-827ccb0eea8a706c4c34a16891f84e7b")
|
||||
os.environ.setdefault("MAIN_MODEL_NAME", "Qwen2.5")
|
||||
|
||||
os.environ.setdefault("EMBEDDING_HTTP_URL", "http://localhost:8023/v1/embeddings")
|
||||
os.environ.setdefault("EMBEDDING_HTTP_KEY", "sk-local-827ccb0eea8a706c4c34a16891f84e7b")
|
||||
os.environ.setdefault("EMBEDDING_MODEL_NAME", "Qwen3-Embedding")
|
||||
# 读取
|
||||
VLLM_HTTP_KEY = os.environ.get("VLLM_HTTP_KEY")
|
||||
VLLM_HTTP_URL = os.environ.get("VLLM_HTTP_URL")
|
||||
VLLM_CHAT_MODEL_NAME = os.environ.get("VLLM_CHAT_MODEL_NAME")
|
||||
MAIN_HTTP_URL = os.environ.get("MAIN_HTTP_URL")
|
||||
MAIN_HTTP_KEY = os.environ.get("MAIN_HTTP_KEY")
|
||||
MAIN_MODEL_NAME = os.environ.get("MAIN_MODEL_NAME")
|
||||
EMBEDDING_HTTP_URL = os.environ.get("EMBEDDING_HTTP_URL")
|
||||
EMBEDDING_HTTP_KEY = os.environ.get("EMBEDDING_HTTP_KEY")
|
||||
EMBEDDING_MODEL_NAME = os.environ.get("EMBEDDING_MODEL_NAME")
|
||||
|
||||
test_prompt = "你好,你是谁,有什么功能,中文回复"
|
||||
print("--- 普通对话 ---")
|
||||
|
||||
text = asyncio.run(call_llm(test_prompt))
|
||||
print(text)
|
||||
# print("--- 普通对话 ---")
|
||||
# text = asyncio.run(call_llm(test_prompt))
|
||||
# print(text)
|
||||
|
||||
print("--- 对话结束 ---\n")
|
||||
# print("--- 流式对话 ---")
|
||||
# async def _test_stream():
|
||||
# try:
|
||||
# async for chunk in call_llm_stream(test_prompt):
|
||||
# print(chunk, end="", flush=True)
|
||||
# except Exception as e:
|
||||
# print(f"_test_stream[ERROR]: {e}")
|
||||
# asyncio.run(_test_stream())
|
||||
|
||||
print("--- 流式对话 ---")
|
||||
print("--- embedding ---")
|
||||
async def _test_embedding():
|
||||
res = await get_embedding(test_prompt)
|
||||
print(res)
|
||||
|
||||
async def _test_stream():
|
||||
try:
|
||||
async for chunk in call_llm_stream(test_prompt):
|
||||
print(chunk, end="", flush=True)
|
||||
print("\n--- 流式结束 ---")
|
||||
except Exception as e:
|
||||
print(f"_test_stream[ERROR]: {e}")
|
||||
asyncio.run(_test_stream())
|
||||
asyncio.run(_test_embedding())
|
||||
|
|
|
|||
|
|
@ -0,0 +1,49 @@
|
|||
import pytest
|
||||
import asyncio
|
||||
|
||||
from sample_ai.nodes import ChunkDocumentsNode
|
||||
from sample_ai.pocketflow import AsyncFlow
|
||||
|
||||
def test_chunk():
|
||||
|
||||
async def _do():
|
||||
shared = {
|
||||
"texts": [
|
||||
# PocketFlow framework
|
||||
"""Pocket Flow is a 100-line minimalist LLM framework
|
||||
Lightweight: Just 100 lines. Zero bloat, zero dependencies, zero vendor lock-in.
|
||||
Expressive: Everything you love—(Multi-)Agents, Workflow, RAG, and more.
|
||||
Agentic Coding: Let AI Agents (e.g., Cursor AI) build Agents—10x productivity boost!
|
||||
To install, pip install pocketflow or just copy the source code (only 100 lines).""",
|
||||
# Fictional medical device
|
||||
"""NeurAlign M7 is a revolutionary non-invasive neural alignment device.
|
||||
Targeted magnetic resonance technology increases neuroplasticity in specific brain regions.
|
||||
Clinical trials showed 72% improvement in PTSD treatment outcomes.
|
||||
Developed by Cortex Medical in 2024 as an adjunct to standard cognitive therapy.
|
||||
Portable design allows for in-home use with remote practitioner monitoring.""",
|
||||
# Made-up historical event
|
||||
"""The Velvet Revolution of Caldonia (1967-1968) ended Generalissimo Verak's 40-year rule.
|
||||
Led by poet Eliza Markovian through underground literary societies.
|
||||
Culminated in the Great Silence Protest with 300,000 silent protesters.
|
||||
First democratic elections held in March 1968 with 94% voter turnout.
|
||||
Became a model for non-violent political transitions in neighboring regions.""",
|
||||
# Fictional technology
|
||||
"""Q-Mesh is QuantumLeap Technologies' instantaneous data synchronization protocol.
|
||||
Utilizes directed acyclic graph consensus for 500,000 transactions per second.
|
||||
Consumes 95% less energy than traditional blockchain systems.
|
||||
Adopted by three central banks for secure financial data transfer.
|
||||
Released in February 2024 after five years of development in stealth mode.""",
|
||||
# Made-up scientific research
|
||||
"""Harlow Institute's Mycelium Strain HI-271 removes 99.7% of PFAS from contaminated soil.
|
||||
Engineered fungi create symbiotic relationships with native soil bacteria.
|
||||
Breaks down "forever chemicals" into non-toxic compounds within 60 days.
|
||||
Field tests successfully remediated previously permanently contaminated industrial sites.
|
||||
Deployment costs 80% less than traditional chemical extraction methods.""",
|
||||
]
|
||||
}
|
||||
|
||||
chunk_node = ChunkDocumentsNode()
|
||||
flow = AsyncFlow(chunk_node)
|
||||
await flow.run_async(shared)
|
||||
|
||||
asyncio.run(_do())
|
||||
|
|
@ -0,0 +1,14 @@
|
|||
import pytest
|
||||
|
||||
from sample_ai.utils import get_embedding, call_llm
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedidng():
|
||||
res = await get_embedding("你好")
|
||||
print(res)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat():
|
||||
res = await call_llm("你好")
|
||||
print(res)
|
||||
Loading…
Reference in New Issue