核心代码展示
展示两个项目的关键实现
DeepNovel_RAG - Token优化实现
class TokenEstimator:
"""Token消耗估算和优化"""
def __init__(self):
self.cache = {}
self.template_cache = {}
def estimate_tokens(self, text: str) -> int:
"""估算Token消耗"""
return len(text) // 4
def optimize_prompt(self, prompt: str) -> str:
"""优化Prompt，减少Token消耗"""
optimized = compress_template(prompt)
cache_key = hash(optimized)
if cache_key in self.cache:
return self.cache[cache_key]
self.cache[cache_key] = optimized
return optimized
xiyou_novel - RAG引擎实现
class RAGEngine:
"""RAG引擎核心"""
def __init__(self, model_provider, vector_store, rerank_client):
self.model_provider = model_provider
self.vector_store = vector_store
self.rerank_client = rerank_client
self.query_optimizer = QueryOptimizer()
def query(self, question: str) -> dict:
"""执行RAG查询"""
optimized_query = self.query_optimizer.optimize(question)
results = self.vector_store.search(optimized_query, top_k=20)
reranked = self.rerank_client.rerank(optimized_query, results)
context = format_context(reranked[:5])
answer = self.model_provider.generate(question, context)
return {
"answer": answer,
"sources": reranked[:5]
}
xiyou_novel - Style-RAG实现
class StyleAnnotator:
"""风格标注器"""
def __init__(self, model_provider):
self.model_provider = model_provider
self.style_dimensions = [
"语言风格", "叙事视角", "修辞手法",
"节奏控制", "人物对话"
]
def annotate(self, text: str) -> dict:
"""标注文本风格"""
prompt = f"""
分析以下文本的写作风格，从以下维度进行标注：
{', '.join(self.style_dimensions)}
文本：
{text}
输出JSON格式的分析结果。
"""
response = self.model_provider.generate(prompt)
return json.loads(response)
def apply_style(self, text: str, style: dict) -> str:
"""应用风格到文本"""
prompt = f"""
请将以下文本按照指定的风格进行改写：
风格要求：
{json.dumps(style, ensure_ascii=False)}
文本：
{text}
"""
return self.model_provider.generate(prompt)