tpoisonooo commited on
Commit
7b67ad7
Β·
1 Parent(s): a7f6abf

feat(examples): support siliconcloud free API

Browse files
README.md CHANGED
@@ -629,6 +629,7 @@ def extract_queries(file_path):
629
  β”‚ β”œβ”€β”€ lightrag_ollama_demo.py
630
  β”‚ β”œβ”€β”€ lightrag_openai_compatible_demo.py
631
  β”‚ β”œβ”€β”€ lightrag_openai_demo.py
 
632
  β”‚ └── vram_management_demo.py
633
  β”œβ”€β”€ lightrag
634
  β”‚ β”œβ”€β”€ __init__.py
 
629
  β”‚ β”œβ”€β”€ lightrag_ollama_demo.py
630
  β”‚ β”œβ”€β”€ lightrag_openai_compatible_demo.py
631
  β”‚ β”œβ”€β”€ lightrag_openai_demo.py
632
+ β”‚ β”œβ”€β”€ lightrag_siliconcloud_demo.py
633
  β”‚ └── vram_management_demo.py
634
  β”œβ”€β”€ lightrag
635
  β”‚ β”œβ”€β”€ __init__.py
examples/lightrag_siliconcloud_demo.py ADDED
@@ -0,0 +1,79 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import asyncio
3
+ from lightrag import LightRAG, QueryParam
4
+ from lightrag.llm import openai_complete_if_cache, siliconcloud_embedding
5
+ from lightrag.utils import EmbeddingFunc
6
+ import numpy as np
7
+
8
+ WORKING_DIR = "./dickens"
9
+
10
+ if not os.path.exists(WORKING_DIR):
11
+ os.mkdir(WORKING_DIR)
12
+
13
+
14
+ async def llm_model_func(
15
+ prompt, system_prompt=None, history_messages=[], **kwargs
16
+ ) -> str:
17
+ return await openai_complete_if_cache(
18
+ "Qwen/Qwen2.5-7B-Instruct",
19
+ prompt,
20
+ system_prompt=system_prompt,
21
+ history_messages=history_messages,
22
+ api_key=os.getenv("UPSTAGE_API_KEY"),
23
+ base_url="https://api.siliconflow.cn/v1/",
24
+ **kwargs,
25
+ )
26
+
27
+
28
+ async def embedding_func(texts: list[str]) -> np.ndarray:
29
+ return await siliconcloud_embedding(
30
+ texts,
31
+ model="netease-youdao/bce-embedding-base_v1",
32
+ api_key=os.getenv("UPSTAGE_API_KEY"),
33
+ max_token_size=int(512 * 1.5)
34
+ )
35
+
36
+
37
+ # function test
38
+ async def test_funcs():
39
+ result = await llm_model_func("How are you?")
40
+ print("llm_model_func: ", result)
41
+
42
+ result = await embedding_func(["How are you?"])
43
+ print("embedding_func: ", result)
44
+
45
+
46
+ asyncio.run(test_funcs())
47
+
48
+
49
+ rag = LightRAG(
50
+ working_dir=WORKING_DIR,
51
+ llm_model_func=llm_model_func,
52
+ embedding_func=EmbeddingFunc(
53
+ embedding_dim=768, max_token_size=512, func=embedding_func
54
+ ),
55
+ )
56
+
57
+
58
+ with open("./book.txt") as f:
59
+ rag.insert(f.read())
60
+
61
+ # Perform naive search
62
+ print(
63
+ rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
64
+ )
65
+
66
+ # Perform local search
67
+ print(
68
+ rag.query("What are the top themes in this story?", param=QueryParam(mode="local"))
69
+ )
70
+
71
+ # Perform global search
72
+ print(
73
+ rag.query("What are the top themes in this story?", param=QueryParam(mode="global"))
74
+ )
75
+
76
+ # Perform hybrid search
77
+ print(
78
+ rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid"))
79
+ )
lightrag/llm.py CHANGED
@@ -2,8 +2,11 @@ import os
2
  import copy
3
  import json
4
  import aioboto3
 
5
  import numpy as np
6
  import ollama
 
 
7
  from openai import AsyncOpenAI, APIConnectionError, RateLimitError, Timeout
8
  from tenacity import (
9
  retry,
@@ -312,7 +315,7 @@ async def ollama_model_complete(
312
  @wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8192)
313
  @retry(
314
  stop=stop_after_attempt(3),
315
- wait=wait_exponential(multiplier=1, min=4, max=10),
316
  retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
317
  )
318
  async def openai_embedding(
@@ -332,6 +335,49 @@ async def openai_embedding(
332
  )
333
  return np.array([dp.embedding for dp in response.data])
334
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
335
 
336
  # @wrap_embedding_func_with_attrs(embedding_dim=1024, max_token_size=8192)
337
  # @retry(
 
2
  import copy
3
  import json
4
  import aioboto3
5
+ import aiohttp
6
  import numpy as np
7
  import ollama
8
+ import base64
9
+ import struct
10
  from openai import AsyncOpenAI, APIConnectionError, RateLimitError, Timeout
11
  from tenacity import (
12
  retry,
 
315
  @wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8192)
316
  @retry(
317
  stop=stop_after_attempt(3),
318
+ wait=wait_exponential(multiplier=1, min=4, max=60),
319
  retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
320
  )
321
  async def openai_embedding(
 
335
  )
336
  return np.array([dp.embedding for dp in response.data])
337
 
338
+ @retry(
339
+ stop=stop_after_attempt(3),
340
+ wait=wait_exponential(multiplier=1, min=4, max=60),
341
+ retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
342
+ )
343
+ async def siliconcloud_embedding(
344
+ texts: list[str],
345
+ model: str = "netease-youdao/bce-embedding-base_v1",
346
+ base_url: str = "https://api.siliconflow.cn/v1/embeddings",
347
+ max_token_size: int = 512,
348
+ api_key: str = None,
349
+ ) -> np.ndarray:
350
+ if api_key and not api_key.startswith('Bearer '):
351
+ api_key = 'Bearer ' + api_key
352
+
353
+ headers = {
354
+ "Authorization": api_key,
355
+ "Content-Type": "application/json"
356
+ }
357
+
358
+ truncate_texts = [text[0:max_token_size] for text in texts]
359
+
360
+ payload = {
361
+ "model": model,
362
+ "input": truncate_texts,
363
+ "encoding_format": "base64"
364
+ }
365
+
366
+ base64_strings = []
367
+ async with aiohttp.ClientSession() as session:
368
+ async with session.post(base_url, headers=headers, json=payload) as response:
369
+ content = await response.json()
370
+ if 'code' in content:
371
+ raise ValueError(content)
372
+ base64_strings = [item['embedding'] for item in content['data']]
373
+
374
+ embeddings = []
375
+ for string in base64_strings:
376
+ decode_bytes = base64.b64decode(string)
377
+ n = len(decode_bytes) // 4
378
+ float_array = struct.unpack('<' + 'f' * n, decode_bytes)
379
+ embeddings.append(float_array)
380
+ return np.array(embeddings)
381
 
382
  # @wrap_embedding_func_with_attrs(embedding_dim=1024, max_token_size=8192)
383
  # @retry(
requirements.txt CHANGED
@@ -11,4 +11,5 @@ tiktoken
11
  torch
12
  transformers
13
  xxhash
14
- pyvis
 
 
11
  torch
12
  transformers
13
  xxhash
14
+ pyvis
15
+ aiohttp