Pythonllama_qa.py
Copy
1
import openai
2
 
3
# Ollama serves Llama on OpenAI-compatible API
4
# Setup: ollama pull llama3.2:3b && ollama serve
5
client = openai.OpenAI(
6
    base_url="http://localhost:11434/v1",
7
    api_key="ollama",
8
)
9
 
10
SYSTEM_PROMPT = """You are a senior code reviewer. Review code for:
11
1. Security vulnerabilities
12
2. Performance issues
13
3. Code quality and readability
14
Be concise and actionable."""
15
 
16
CODE_TO_REVIEW = '''
17
def get_user(user_id):
18
    query = f"SELECT * FROM users WHERE id = {user_id}"
19
    conn = db.connect()
20
    result = conn.execute(query)
21
    return result.fetchall()
22
'''
23
 
24
def review_code(code: str, model: str = "llama3.2:3b") -> str:
25
    response = client.chat.completions.create(
26
        model=model,
27
        messages=[
28
            {"role": "system", "content": SYSTEM_PROMPT},
29
            {"role": "user", "content": f"Review this code:\n\n{code}"},
30
        ],
31
        max_tokens=500,
32
        temperature=0.1,
33
    )
34
    return response.choices[0].message.content
35
 
36
# Document Q&A with chunking
37
def chunk_document(text: str, chunk_size: int = 500) -> list[str]:
38
    words = text.split()
39
    chunks = []
40
    for i in range(0, len(words), chunk_size):
41
        chunks.append(" ".join(words[i:i + chunk_size]))
42
    return chunks
43
 
44
def qa_pipeline(document: str, question: str, model: str = "llama3.2:3b") -> str:
45
    """Simple Q&A: find relevant chunk, then answer."""
46
    chunks = chunk_document(document)
47
    # In production: use embeddings for semantic search
48
    # Here: keyword-based retrieval
49
    relevant = [c for c in chunks if any(w.lower() in c.lower() 
50
                for w in question.lower().split() if len(w) > 4)]
51
    context = "\n\n".join(relevant[:2]) if relevant else chunks[0]
52
    
53
    response = client.chat.completions.create(
54
        model=model,
55
        messages=[
56
            {"role": "system", "content": "Answer the question using only the provided context."},
57
            {"role": "user", "content": f"Context: {context}\n\nQuestion: {question}"},
58
        ],
59
        max_tokens=200,
60
        temperature=0,
61
    )
62
    return response.choices[0].message.content
63
 
64
print("=== Code Review ===")
65
try:
66
    review = review_code(CODE_TO_REVIEW)
67
    print(review)
68
except Exception as e:
69
    print(f"Error (is ollama running?): {e}")
70